Fixing pytorch lightning dependency

flatironinstitute · Mar 8, 2022 · 05aec99 · 05aec99
1 parent ab48e4f
commit 05aec99
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/deepblast/tests/test_trainer.py b/deepblast/tests/test_trainer.py
@@ -1,6 +1,7 @@
 import os
 import shutil
 import unittest
+import torch
 from deepblast.trainer import LightningAligner
 from deepblast.utils import get_data_path
 from deepblast.sim import hmm_alignments
@@ -39,7 +40,7 @@ def tearDown(self):
         if os.path.exists('valid.txt'):
             os.remove('valid.txt')
 
-    @unittest.skip("Can only run with GPU")
+    @unittest.skipUnless(torch.cuda.is_available(), 'No GPU was detected')
     def test_trainer_sim(self):
         output_dir = 'output'
         args = [
@@ -51,7 +52,7 @@ def test_trainer_sim(self):
             '--batch-size', '3',
             '--num-workers', '1',
             '--learning-rate', '1e-4',
-            '--clip-ends', 'False',
+            # '--clip-ends', 'False',
             '--visualization-fraction', '0.5',
             '--gpus', '1'
         ]
@@ -71,7 +72,7 @@ def test_trainer_sim(self):
         )
         trainer.fit(model)
 
-    @unittest.skip("Can only run with GPU")
+    @unittest.skipUnless(torch.cuda.is_available(), 'No GPU was detected')
     def test_trainer_struct(self):
         output_dir = 'output'
         args = [
@@ -83,7 +84,7 @@ def test_trainer_struct(self):
             '--batch-size', '1',
             '--num-workers', '1',
             '--learning-rate', '1e-4',
-            '--clip-ends', 'False',
+            # '--clip-ends', 'False',
             '--visualization-fraction', '1',
             '--gpus', '1'
         ]

diff --git a/deepblast/trainer.py b/deepblast/trainer.py
@@ -25,8 +25,9 @@ class LightningAligner(pl.LightningModule):
 
     def __init__(self, args):
         super(LightningAligner, self).__init__()
+        self._hparams = args
         self.tokenizer = UniprotTokenizer(pad_ends=False)
-        self.hparams = args
+        #self.hparams = args
         self.initialize_aligner()
         if self.hparams.loss == 'sse':
             self.loss_func = SoftAlignmentLoss()