legacy bug fix for the trainer assuming there exists a model.mode attribute when its not really required (#545)

Co-authored-by: John Wu <johnwu3@sunlab-serv-03.cs.illinois.edu>

Author: jhnwu3

pixi.lock

@@ -30,6 +30,8 @@ def __init__(self, dataset: SampleDataset):
         # used to query the device of the model
         self._dummy_param = nn.Parameter(torch.empty(0))
 
+        self.mode = None  # legacy API for backward compatibility with the trainer.
+
     @property
     def device(self) -> torch.device:
         """
@@ -50,9 +52,9 @@ def get_output_size(self) -> int:
         Returns:
             int: The output size of the model.
         """
-        assert len(self.label_keys) == 1, (
-            "Only one label key is supported if get_output_size is called"
-        )
+        assert (
+            len(self.label_keys) == 1
+        ), "Only one label key is supported if get_output_size is called"
         output_size = self.dataset.output_processors[self.label_keys[0]].size()
         return output_size
 
@@ -69,9 +71,9 @@ def get_loss_function(self) -> Callable:
         Returns:
             Callable: The default loss function.
         """
-        assert len(self.label_keys) == 1, (
-            "Only one label key is supported if get_loss_function is called"
-        )
+        assert (
+            len(self.label_keys) == 1
+        ), "Only one label key is supported if get_loss_function is called"
         label_key = self.label_keys[0]
         mode = self.dataset.output_schema[label_key]
         if mode == "binary":
@@ -106,9 +108,9 @@ def prepare_y_prob(self, logits: torch.Tensor) -> torch.Tensor:
         Returns:
             torch.Tensor: The predicted probability tensor.
         """
-        assert len(self.label_keys) == 1, (
-            "Only one label key is supported if get_loss_function is called"
-        )
+        assert (
+            len(self.label_keys) == 1
+        ), "Only one label key is supported if get_loss_function is called"
         label_key = self.label_keys[0]
         mode = self.dataset.output_schema[label_key]
         if mode in ["binary"]: