def training_step(self, batch, batch_idx):
images, masks, _, _ = batch
images, masks = Utils.preprocessing(images, masks, self.WL, self.WW)
images, masks = Utils.do_train_augmentations(images, masks,
self.gaussian_noise_std,
self.device, self.ra,
self.rf)
y_hat = self(images)
# loss dim is [batch, 1, img_x, img_y]
# need to get rid of the second dimension so
# size matches with mask
loss = self.loss(y_hat[:, 0, :, :], masks)
# Logs
#tensorboard_logs = {'train_loss': loss}
return {'loss': loss} #, 'log': tensorboard_logs}
