【Pytorch Lightning × Google Colaboratory】Tnesorboardでデータを可視化する方法

class TrainNet(pl.LightningModule):

    @pl.data_loader
    def train_dataloader(self):
        return torch.utils.data.DataLoader(train, self.batch_size,shuffle=True)
    
    def training_step(self, batch, batch_nb):
        x, t = batch
        y = self.forward(x)
        loss = self.lossfun(y, t)
　　　　　##### 追加 ######
        tensorboard_logs = {'train_loss': loss}
        results = {'loss': loss, 'log': tensorboard_logs}
　　　　　################
        return results

class ValidationNet(pl.LightningModule):

    @pl.data_loader
    def val_dataloader(self):
        return torch.utils.data.DataLoader(val, self.batch_size)

    def validation_step(self, batch, batch_nb):
        x, t = batch
        y = self.forward(x)
        loss = self.lossfun(y, t)
        y_label = torch.argmax(y, dim=1)
        acc = torch.sum(t == y_label) * 1.0 / len(t)
        results = {'val_loss': loss, 'val_acc': acc}
        return results

    def validation_end(self, outputs):
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['val_acc'] for x in outputs]).mean()
　　　　　##### 追加 ######
        tensorboard_logs = {'avg_acc': avg_acc}
        results =  {'val_loss': avg_loss, 'val_acc': avg_acc, 'log': tensorboard_logs}
　　　　　################
        return results

%load_ext tensorboard
%tensorboard --logdir ./lightning_logs/フォルダ名