[Misc] disable cast_forward_inputs (#460)

2025-05-31 22:06:29 -07:00
parent b2ebaaf865
commit dcf7738cbc
1 changed files with 3 additions and 1 deletions
@@ -102,10 +102,12 @@ def load_fsdp_model(
    output_dtype: Optional[torch.dtype] = None,
 ) -> torch.nn.Module:

+    # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
+    # manually casting the inputs to the model
    mp_policy = MixedPrecisionPolicy(param_dtype,
                                     reduce_dtype,
                                     output_dtype,
-                                     cast_forward_inputs=True)
+                                     cast_forward_inputs=False)

    with set_default_dtype(default_dtype), torch.device("meta"):
        model = model_cls(**init_params)