cropped loss now a torch module for checkpointing
import torch
+from torch import nn
-class CroppedLoss:
+class CroppedLoss(nn.Module):
"""Compute Loss after averaging predictions across time.
Assumes predictions are in shape:
n_batch size x n_classes x n_predictions (in time)"""
def __init__(self, loss_function):
+ super().__init__()
self.loss_function = loss_function
- def __call__(self, preds, targets):
+ def forward(self, preds, targets):
avg_preds = torch.mean(preds, dim=2)
avg_preds = avg_preds.squeeze(dim=1)
return self.loss_function(avg_preds, targets)
