Merge pull request #234 from google-ai-edge/pytorchoptional

Show pytorch adapter even if torch in not imported.
google-ai-edge · Nov 4, 2024 · a48865d · a48865d
2 parents 6aa67b6 + d4cb9db
commit a48865d
Showing 1 changed file with 9 additions and 18 deletions.
diff --git a/src/server/package/src/model_explorer/extension_manager.py b/src/server/package/src/model_explorer/extension_manager.py
@@ -19,11 +19,6 @@
 from importlib import import_module
 from typing import Any, Dict, Union
 
-try:
-  import torch
-except ImportError:
-  torch = None
-
 from .adapter_runner import AdapterRunner
 from .consts import MODULE_NAME
 from .extension_class_processor import ExtensionClassProcessor
@@ -34,19 +29,15 @@
 
 
 class ExtensionManager(object, metaclass=Singleton):
-  BUILTIN_ADAPTER_MODULES: list[str] = (
-      [
-          '.builtin_tflite_flatbuffer_adapter',
-          '.builtin_tflite_mlir_adapter',
-          '.builtin_tf_mlir_adapter',
-          '.builtin_tf_direct_adapter',
-          '.builtin_graphdef_adapter',
-      ]
-      + (['.builtin_pytorch_exportedprogram_adapter'] if torch else [])
-      + [
-          '.builtin_mlir_adapter',
-      ]
-  )
+  BUILTIN_ADAPTER_MODULES: list[str] = [
+      '.builtin_tflite_flatbuffer_adapter',
+      '.builtin_tflite_mlir_adapter',
+      '.builtin_tf_mlir_adapter',
+      '.builtin_tf_direct_adapter',
+      '.builtin_graphdef_adapter',
+      '.builtin_pytorch_exportedprogram_adapter',
+      '.builtin_mlir_adapter',
+  ]
 
   CACHED_REGISTERED_EXTENSIONS: Dict[str, RegisteredExtension] = {}