Fix classification labels for batch size 1

lukasugar · lukasugar · commit 964e8d25daa7 · 2026-05-18T17:56:40.000+02:00
diff --git a/asparagus/modules/lightning_modules/clsreg_module.py b/asparagus/modules/lightning_modules/clsreg_module.py
@@ -188,7 +188,7 @@ def configure_metrics(self, prefix: str):
         )
 
     def on_before_batch_transfer(self, batch, dataloader_idx):
-        batch["CLSREG_label"] = batch["CLSREG_label"].squeeze().long()
+        batch["CLSREG_label"] = batch["CLSREG_label"].view(-1).long()
         return batch
 
     def on_test_batch_end(self, outputs, batch, batch_idx, dataloader_idx=0):

Original file line number	Diff line number	Diff line change
`@@ -188,7 +188,7 @@ def configure_metrics(self, prefix: str):`
`188`	`188`	`)`
`189`	`189`
`190`	`190`	`def on_before_batch_transfer(self, batch, dataloader_idx):`
`191`		`- batch["CLSREG_label"] = batch["CLSREG_label"].squeeze().long()`
	`191`	`+ batch["CLSREG_label"] = batch["CLSREG_label"].view(-1).long()`
`192`	`192`	`return batch`
`193`	`193`
`194`	`194`	`def on_test_batch_end(self, outputs, batch, batch_idx, dataloader_idx=0):`