[pyTorch] Enable the model to change precision between iterations (#414)

* Enable the model to be change precision between iterations
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Add test
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix for the test
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

tests/pytorch/test_sanity.py

@@ -788,3 +788,16 @@ def test_gpt_cuda_graph(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_
     )
 
     _test_sanity_e2e_cuda_graph(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+def test_model_multiple_cast():
+    a = torch.zeros((16,16)).cuda()
+    m = Linear(16,32)
+
+    y = m(a)
+    assert y.dtype == torch.float32
+
+    m.half()
+    a = a.half()
+
+    y2 = m(a)
+    assert y2.dtype == torch.float16

transformer_engine/pytorch/module/base.py

@@ -445,8 +445,7 @@ class TransformerEngineBaseModule(torch.nn.Module, ABC):
             return
 
         # All checks after this have already been performed once, thus skip
-        # We assume that user doesn't change input types across iterations
-        if hasattr(self, "activation_dtype"):
+        if hasattr(self, "activation_dtype") and self.activation_dtype == inp.dtype:
             return
 
         dtype = inp.dtype