Moved mixed data

Fixed code to one hot new mask tensors

src/data_collection/data_loader.py

@@ -75,6 +75,7 @@ def sample(self, batch_size: int, time_steps: int) -> Tuple[torch.Tensor, torch.
         states_tensor  = states_tensor .reshape((-1, 12, 128, 128))
 
         masks_tensor = torch.from_numpy(np.array(masks))[:, :self.num_obj]
-        masks_tensor = F.one_hot(masks_tensor, num_classes=self.num_obj).float()[1:]  # get rid of background
+        masks_tensor = F.one_hot(masks_tensor.long(), num_classes=self.num_obj + 1).float()[1:]  # get rid of background [B, H, W, O]
+        masks_tensor = masks_tensor.permute(0, 1, 4, 2, 3)  # [B, O, H, W]
 
         return states_tensor, object_bounding_boxes_tensor, masks_tensor, torch.from_numpy(np.array(actions))

src/scripts/train_prediction.py

@@ -17,7 +17,7 @@ def train(config: DictConfig, batch_size: int = 4, t_steps: int = 1, num_obj: in
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     print(f"Using device: {device}")
 
-    data_loader = DataLoader("SimpleTestDataSmall", num_obj)
+    data_loader = DataLoader("SimpleTestData", num_obj)
 
     feature_extract = FeatureExtractor(num_objects=num_obj).to(device)
     predictor = Predictor(num_layers=1, time_steps=t_steps).to(device)