[fix] gpt2 neuron support handler and ci (#1229)

deepjavalibrary · Oct 27, 2023 · 5e25d61 · 5e25d61
1 parent 7ed7270
commit 5e25d61
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/engines/python/setup/djl_python/transformers_neuronx.py b/engines/python/setup/djl_python/transformers_neuronx.py
@@ -180,12 +180,7 @@ def load_model(self, model_type):
         # TODO: workaround on Neuron Compiler bug for SM
         path = os.getcwd()
         os.chdir("/tmp")
-        if model_type == "gpt2":
-            self.model._load_compiled_artifacts(load_path)
-            self.model.to_neuron()
-            self.model._save_compiled_artifacts(load_path)
-        else:
-            self.model.to_neuron()
+        self.model.to_neuron()
         os.chdir(path)
         elapsed = time.time() - start
         logging.info(