[PyTorch] Fix for type checking failure on custom callables (#790)

fix type checking in checkpointing to assume that there must be TE modules in custom callables
Signed-off-by: Alp Dener <adener@nvidia.com>

transformer_engine/pytorch/distributed.py

@@ -521,9 +521,12 @@ def has_te_modules(network):
         for module in network.modules():
             if any(isinstance(module, te_class) for te_class in te_classes_list):
                 return True
-
         return False
 
+    # Cannot check for TE modules inside a custom class/callable that's not a torch.nn.Module,
+    # so just assume that it has TE modules just to be safe.
+    return True
+
 
 def checkpoint(
     function: Callable,