Make sure to register adversarial module with Fabric.

arnor-sigurdsson · arnor-sigurdsson · commit 58fe80564b9b · 2026-04-16T10:53:39.000-06:00
diff --git a/src/eir/train_utils/adversarial.py b/src/eir/train_utils/adversarial.py
@@ -155,7 +155,7 @@ def hook_add_adversarial_losses(
                     adv_config.projection_lcl_residual_blocks
                 ),
             )
-            module = module.to(device=device)
+            module = experiment.fabric.setup_module(module)
             adversarial_modules[adv_config.name] = module
 
             logger.debug(

Original file line number	Diff line number	Diff line change
`@@ -155,7 +155,7 @@ def hook_add_adversarial_losses(`
`155`	`155`	`adv_config.projection_lcl_residual_blocks`
`156`	`156`	`),`
`157`	`157`	`)`
`158`		`- module = module.to(device=device)`
	`158`	`+ module = experiment.fabric.setup_module(module)`
`159`	`159`	`adversarial_modules[adv_config.name] = module`
`160`	`160`
`161`	`161`	`logger.debug(`