training fix for test

ricsinaruto · Mar 28, 2023 · 553cff7 · 553cff7
1 parent ff4fd8a
commit 553cff7
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -16,4 +16,5 @@ transformers
 PyWavelets
 hyperopt
 pyriemann
-xgboost
+xgboost
+sails
diff --git a/training.py b/training.py
@@ -80,10 +80,7 @@ def __init__(self, args, dataset=None):
             #self.args.dataset = self.dataset
         else:
             self.model_path = os.path.join(self.args.result_dir, 'model.pt')
-            if args.from_pretrained:
-                self.model = args.model.from_pretrained(args)
-            else:
-                self.model = args.model(args)
+            self.model = args.model(args)
 
             try:
                 self.model = self.model.cuda()

diff --git a/transformers_quantized.py b/transformers_quantized.py
@@ -112,11 +112,15 @@ class TransformerQuantizedPretrained(TransformerQuantized):
     '''
     Same as TransformerQuantized, but uses the pretrained GPT2 model
     '''
+    def __init__(self, args):
+        super().__init__(args.gpt2_config)
+
     @classmethod
     def from_pretrained(cls, args):
         return super().from_pretrained('gpt2',
                                        args,
-                                       config=args.gpt2_config)
+                                       config=args.gpt2_config,
+                                       cache_dir=args.result_dir)
 
 
 class TransformerQuantizedConcatEmb(TransformerQuantized):