Move SWA Model to AutoUnit.device

Summary: Update torchtnt auto_unit to use self.device for the EMA / SWA model, which may be set from environment in the superclass init. This enables model evaluation in GPU. Differential Revision: D64206735
pytorch · Oct 11, 2024 · ca5d580 · ca5d580
1 parent b9e7c1f
commit ca5d580
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchtnt/framework/auto_unit.py b/torchtnt/framework/auto_unit.py
@@ -512,7 +512,7 @@ def __init__(
 
             self.swa_model = AveragedModel(
                 module_for_swa,
-                device=device,
+                device=self.device,
                 use_buffers=swa_params.use_buffers,
                 averaging_method=swa_params.averaging_method,
                 ema_decay=swa_params.ema_decay,