Example training function with tensorboardcolab

train.py

def train(model, train_loader, device, optimizer, log_interval, epoch, globaliter, tb):
  """
  Example training function for PyTorch recording to tensorboardcolab. 
  """

  model.train()
  for batch_idx, (data, target) in enumerate(train_loader):

    globaliter += 1
    data, target = data.to(device), target.to(device)

    optimizer.zero_grad()
    predictions = model(data)

    loss = F.nll_loss(predictions, target)
    loss.backward()
    optimizer.step()

    if batch_idx % log_interval == 0:
      print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

      # This is where I'm recording to Tensorboard
      tb.save_value('Train Loss', 'train_loss', self.globaliter, loss.item())