Switch AC UNet channels

htem · Nov 29, 2023 · e25ceea · e25ceea
1 parent 20327ca
commit e25ceea
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 1 deletion.
diff --git a/log/log_aclsd/events.out.tfevents.1701276331.lee-htem-gpu0 b/log/log_aclsd/events.out.tfevents.1701276331.lee-htem-gpu0
diff --git a/log/log_mtlsd/events.out.tfevents.1701276331.lee-htem-gpu0 b/log/log_mtlsd/events.out.tfevents.1701276331.lee-htem-gpu0
diff --git a/src/autoseg/train/ACLSDTrain.py b/src/autoseg/train/ACLSDTrain.py
@@ -44,7 +44,7 @@ def aclsd_train(
 
     # initial MTLSD UNet
     unet = setup_unet(downsample_factors=[(2, 2, 2), (2, 2, 2)], padding="same")
-    unet_ac = setup_unet(downsample_factors=[(2, 2, 2), (2, 2, 2)], padding="same", num_heads=1)
+    unet_ac = setup_unet(in_channels=10, downsample_factors=[(2, 2, 2), (2, 2, 2)], padding="same", num_heads=1)
 
     mtlsd_model = MTLSDModel(unet=unet, num_fmaps=unet.out_channels)
     mtlsd_loss = Weighted_MSELoss()  # aff_lambda=0)