handle ImportError and OSError in extras.pytorch (#1141)

* handle ImportError and OSError in extras.pytorch Signed-off-by: Niels Bantilan <[email protected]> * isolate exception to torch import Signed-off-by: Niels Bantilan <[email protected]> Signed-off-by: Niels Bantilan <[email protected]>
flyteorg · Aug 25, 2022 · 8e91b47 · 8e91b47
1 parent c4ba45a
commit 8e91b47
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/flytekit/extras/pytorch/__init__.py b/flytekit/extras/pytorch/__init__.py
@@ -11,10 +11,21 @@
 """
 from flytekit.loggers import logger
 
+# TODO: abstract this out so that there's an established pattern for registering plugins
+# that have soft dependencies
 try:
+    # isolate the exception to the torch import
+    import torch
+
+    _torch_installed = True
+except (ImportError, OSError):
+    _torch_installed = False
+
+
+if _torch_installed:
     from .checkpoint import PyTorchCheckpoint, PyTorchCheckpointTransformer
     from .native import PyTorchModuleTransformer, PyTorchTensorTransformer
-except ImportError:
+else:
     logger.info(
         "We won't register PyTorchCheckpointTransformer, PyTorchTensorTransformer, and PyTorchModuleTransformer because torch is not installed."
     )