lint

basetenlabs · Apr 18, 2024 · e6d3a06 · e6d3a06
1 parent 4fee761
commit e6d3a06
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/distilbert/config.yaml b/distilbert/config.yaml
@@ -3,4 +3,4 @@ model_name: DistilBert
 python_version: py310
 requirements_file: ./requirements.txt
 resources:
-  accelerator: T4
+  accelerator: T4
diff --git a/distilbert/model/model.py b/distilbert/model/model.py
@@ -9,7 +9,7 @@ def __init__(self, **kwargs):
     def load(self):
         # Load model here and assign to self._model.
         self.device = (
-            "cuda" if torch.cuda.is_available() else "mps"
+            "cuda" if torch.cuda.is_available() else "cpu"
         )  # the device to load the model onto
 
         self._tokenizer = AutoTokenizer.from_pretrained(