training is running without error

run_sbatch.sbatch

@@ -1,16 +1,13 @@
 #! /bin/bash
 #SBATCH -c 4 #4 #8
-#SBATCH --mem 120G #120G #32G #64G #256G
+#SBATCH --mem 256G #120G #32G #64G #256G
 #SBATCH -p grete:shared
-#SBATCH -t 5:00:00 #6:00:00 #48:00:00 
+#SBATCH -t 48:00:00 #6:00:00 #48:00:00 
 #SBATCH -G A100:1 #V100:1 #2 #A100:1  #gtx1080:2 #v100:1 #H100:1
 #SBATCH --output=/user/muth9/u12095/czii-protein-challenge/slurm/slurm-%j.out 
-#SBATCH -A nim00007 #SBATCH --constraint 80gb 
+#SBATCH -A nim00007 
+#SBATCH --constraint 80gb 
 
 source ~/.bashrc
 conda activate deeplearning
-python /user/muth9/u12095/czii-protein-challenge/inference/run_protein_detection.py \
-    -i /scratch-grete/projects/nim00007/cryo-et/challenge-data/train/static/ExperimentRuns/TS_69_2 -f \
-    -o /mnt/lustre-emmy-hdd/usr/u12095/cryo-et/czii_challenge/detections \
-    -m /mnt/lustre-emmy-hdd/usr/u12095/cryo-et/czii_challenge/models/checkpoints/protein_detection_czii_v2 \
-    -j /mnt/lustre-emmy-hdd/usr/u12095/cryo-et/czii_challenge/training/protein_detection_czii_v2/split-ExperimentRuns.json
+python /user/muth9/u12095/czii-protein-challenge/training/train.py

utils/training/data_loader.py

@@ -27,10 +27,8 @@ def _load_dataset(
     dataset_class=HeatmapDataset,
     n_samples=None,
 ):
-    print(f"in _load_dataset raw_paths {raw_paths}")
 
     if isinstance(raw_paths, str):
-        print(f"in isinstance raw_paths {raw_paths}")
         ds = dataset_class(
         raw_path=raw_paths, raw_key=raw_key, label_path=label_paths, patch_shape=patch_shape,
         raw_transform=raw_transform, transform=transform, eps=eps, sigma=sigma,
@@ -44,7 +42,6 @@ def _load_dataset(
         )
         ds = []
         for i, (raw_path, label_path) in enumerate(zip(raw_paths, label_paths)):
-            print(f"in else raw_path {raw_path}")
 
             dset = dataset_class(
             raw_path=raw_path, raw_key=raw_key, label_path=label_path, patch_shape=patch_shape,

utils/training/heatmap_dataset.py

@@ -50,18 +50,18 @@ def __init__(
     ):
         self.raw_path = raw_path
         self.raw_key = raw_key
-        print(f"raw_path {raw_path}, raw_key {raw_key}")
+        #print(f"raw_path {raw_path}, raw_key {raw_key}")
         self.raw = load_data(raw_path, raw_key)
-        print(f"self.raw {self.raw}")
+        #print(f"self.raw {self.raw}")
 
         self.label_path = label_path
-        print(f"self.label_path {self.label_path}")
+        #print(f"self.label_path {self.label_path}")
 
         self._with_channels = with_channels
         self._with_label_channels = with_label_channels
 
         shape_raw = self.raw.shape[1:] if self._with_channels else self.raw.shape
-        print(f"shape_raw {shape_raw}")
+        #print(f"shape_raw {shape_raw}")
 
         self.shape = shape_raw