def training_step(self, batch):
    self.model.train()
    fwd_data = self.forward(batch)
    self.optimizer.zero_grad()
    losses_dict = self._compute_lossses(batch, fwd_data)
    loss = losses_dict['loss']
    loss.backward()
    self.optimizer.step()
    self._step += 1
    return {k: v.item() for (k, v) in losses_dict.items()}
def training_step(self, batch):
    self.model.train()
    fwd_data = self.forward(batch)
    self.optimizer.zero_grad()
    losses_dict = self._compute_lossses(batch, fwd_data)
    loss = losses_dict['loss']
    loss.backward()
    self.optimizer.step()
    self._step += 1
    return {k: v.item() for (k, v) in losses_dict.items()}