update train_csn.py

xing-lab-pitt · Mar 19, 2024 · b530ab9 · b530ab9
1 parent 297d4be
commit b530ab9
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/livecellx/model_zoo/segmentation/train_csn.py b/livecellx/model_zoo/segmentation/train_csn.py
@@ -156,6 +156,9 @@ def df2dataset(df):
     if args.debug:
         train_df = train_df[:100]
         val_df = val_df[:10]
+    # Get rid of last batch if it's not full to prevent BatchNorm complaints
+    train_df = train_df.iloc[: (len(train_df) // args.batch_size) * args.batch_size]
+    val_df = val_df.iloc[: (len(val_df) // args.batch_size) * args.batch_size]
 
     train_dataset = df2dataset(train_df)
     val_dataset = df2dataset(val_df)