small fix on ssl_label to None if under supervise training.

Author: Kin-Zhang

train.py

@@ -57,7 +57,7 @@ def main(cfg):
     precheck_cfg_valid(cfg)
     pl.seed_everything(cfg.seed, workers=True)
 
-    train_dataset = HDF5Dataset(cfg.train_data, n_frames=cfg.num_frames, ssl_label=cfg.ssl_label)
+    train_dataset = HDF5Dataset(cfg.train_data, n_frames=cfg.num_frames, ssl_label=cfg.get('ssl_label', None))
     train_loader = DataLoader(train_dataset,
                               batch_size=cfg.batch_size,
                               shuffle=True,
@@ -134,7 +134,7 @@ def main(cfg):
         print("Initiating wandb and trainer successfully.  ^V^ ")
         print(f"We will use {cfg.gpus} GPUs to train the model. Check the checkpoints in {output_dir} checkpoints folder.")
         print("Total Train Dataset Size: ", len(train_dataset))
-        if cfg.add_seloss is not None and cfg.loss_fn in ['seflowLoss', 'seflowppLoss']:
+        if cfg.get('add_seloss', None) is not None and cfg.loss_fn in ['seflowLoss', 'seflowppLoss']:
             print(f"Note: We are in **self-supervised** training now. No ground truth label is used.")
             print(f"We will use these loss items in {cfg.loss_fn}: {cfg.add_seloss}")
         print("-"*40+"\n")