Fix evaluation

LukasHedegaard · Dec 14, 2021 · eb4ee76 · eb4ee76
1 parent 948779d
commit eb4ee76
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 5 deletions.
diff --git a/main.py b/main.py
@@ -196,6 +196,7 @@ def main(args):
     for epoch in range(args.start_epoch, args.epochs):
         if args.distributed:
             sampler_train.set_epoch(epoch)
+
         train_stats = train_one_epoch(
             model,
             criterion,

diff --git a/train.py b/train.py
@@ -79,7 +79,7 @@ def train_one_epoch(
         # enc_score_p0 = model(camera_inputs, motion_inputs)
         enc_score_p0 = model(
             torch.cat((camera_inputs, motion_inputs), 2).transpose(1, 2)
-        )
+        ).squeeze(-1)
 
         outputs = {
             "labels_encoder": enc_score_p0,  # [128, 22]
@@ -174,7 +174,7 @@ def evaluate(model, criterion, data_loader, device, logger, args, epoch, nprocs=
 
         enc_score_p0 = model(
             torch.cat((camera_inputs, motion_inputs), 2).transpose(1, 2)
-        )
+        ).squeeze(-1)
         # set_trace()
 
         outputs = {

diff --git a/util/loss.py b/util/loss.py
@@ -44,9 +44,6 @@ def loss_labels(self, input, targets, name):
         """
         target = targets.float()
 
-        if len(input.shape) > 2 and input.shape[-1] == 1:
-            input = input.squeeze(-1)
-
         if self.ignore_index >= 0:
             notice_index = [
                 i for i in range(target.shape[-1]) if i != self.ignore_index