Updates metric tracking so checkpoints are saved correctly
" return {'val_loss': loss, 'val_dice': value}\n",
" \n",
" def validation_epoch_end(self, outputs):\n",
- " val_dice = 0\n",
- " num_items = 0\n",
+ " val_dice, val_loss, num_items = 0, 0, 0\n",
" for output in outputs:\n",
" val_dice += output['val_dice'].sum().item()\n",
+ " val_loss += output['val_loss'].sum().item()\n",
" num_items += len(output['val_dice'])\n",
- " mean_val_dice = val_dice / num_items\n",
- " tensorboard_logs = {'val_dice': mean_val_dice}\n",
+ " mean_val_dice = torch.tensor(val_dice / num_items)\n",
+ " mean_val_loss = torch.tensor(val_loss / num_items)\n",
+ " tensorboard_logs = {'val_dice': mean_val_dice, 'val_loss': mean_val_loss}\n",
" if mean_val_dice > self.best_val_dice:\n",
" self.best_val_dice = mean_val_dice\n",
" self.best_val_epoch = self.current_epoch\n",
