Merge branch 'main' of github.com:xing-lab-pitt/livecellx into main

xing-lab-pitt · Mar 20, 2024 · 3489fbe · 3489fbe
2 parents 0779321 + 7a6b9a3
commit 3489fbe
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 avoid 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)