add gpu cuda

load.py

@@ -40,11 +40,11 @@ class LoadModel:
 
     def load_model(self):
         if self.revision is None:
-            model = GPT2LMHeadModel.from_pretrained(self.path).to(self.device)
+            model = GPT2LMHeadModel.from_pretrained(self.path, device_map="auto")
         else:
             model = GPT2LMHeadModel.from_pretrained(
-                self.path, revision=self.revision
-            ).to(self.device)
+                self.path, revision=self.revision, device_map="auto"
+            )
 
         return model
 

requirements.txt

@@ -13,4 +13,5 @@ scipy
 datetime
 torch
 torchvision
-pyFluidSynth
+pyFluidSynth
+accelerate