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

Author: Victor Bourgin

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,