lint

Author: bolasim

distilbert/config.yaml

@@ -3,4 +3,4 @@ model_name: DistilBert
 python_version: py310
 requirements_file: ./requirements.txt
 resources:
-  accelerator: T4
+  accelerator: T4

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(