fix test

nnaakkaaii · Jul 14, 2024 · f644551 · f644551
1 parent 29c28b8
commit f644551
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/hrdae/conf b/hrdae/conf
diff --git a/hrdae/models/losses/triplet.py b/hrdae/models/losses/triplet.py
@@ -35,7 +35,7 @@ def forward(
     ) -> Tensor:
         num_frames = positive[0].size(1)
 
-        anchor = latent[0].unsqueeze(1).expand(-1, num_frames, -1, -1, -1)
+        anchor = latent[0].expand(-1, num_frames, -1, -1, -1)
         pos_dist = torch.norm(anchor - positive[0], p=2, dim=[2, 3, 4])
         neg_dist = torch.norm(anchor - negative[0], p=2, dim=[2, 3, 4])
 

diff --git a/test/models/test_vr_model.py b/test/models/test_vr_model.py
@@ -112,6 +112,7 @@ def test_vr_model():
         optimizer,
         scheduler,
         criterion,
+        False,
     )
     with TemporaryDirectory() as tempdir:
         model.train(