Remove typecasting of labels

sage/trainer/trainer.py

@@ -190,7 +190,7 @@ def forward(self, batch, mode: str = "train"):
             """
             aug = getattr(self, f"{'train' if mode == 'train' else 'valid'}_transforms")
             batch["brain"] = aug(batch["brain"]).as_tensor()
-            batch["age"] = batch["age"].float()
+            batch["age"] = batch["age"]
             result: dict = self.model(**batch)
             return result
         except RuntimeError as e: