add device_id for non meta models

abhi-mosaic · abhi-mosaic · commit 98f6e1234b94 · 2022-10-13T17:47:34.000-07:00
diff --git a/composer/trainer/dist_strategy.py b/composer/trainer/dist_strategy.py
@@ -238,6 +238,7 @@ def _auto_wrap_policy(module: torch.nn.Module, recurse: bool, unwrapped_params:
                 mixed_precision=mixed_precision,
                 backward_prefetch=backward_prefetch,
                 param_init_fn=_param_init_fn,
+                device_id=torch.cuda.current_device(),
             )
 
             # Activation Checkpointing

Original file line number	Diff line number	Diff line change
`@@ -238,6 +238,7 @@ def _auto_wrap_policy(module: torch.nn.Module, recurse: bool, unwrapped_params:`
`238`	`238`	`mixed_precision=mixed_precision,`
`239`	`239`	`backward_prefetch=backward_prefetch,`
`240`	`240`	`param_init_fn=_param_init_fn,`
	`241`	`+ device_id=torch.cuda.current_device(),`
`241`	`242`	`)`
`242`	`243`
`243`	`244`	`# Activation Checkpointing`