fix: mypy errors

Author: djaniak

embeddings/pipeline/hps_pipeline.py

@@ -39,7 +39,7 @@ class TorchGarbageCollectorOptunaCallback(OptunaCallback):
     def _cleanup() -> None:
         gc.collect()
         if torch.cuda.is_available():
-            torch.cuda.empty_cache()  # type: ignore
+            torch.cuda.empty_cache()
 
     def __call__(self, study: Study, trial: optuna.trial.FrozenTrial) -> None:
         TorchGarbageCollectorOptunaCallback._cleanup()

embeddings/task/lightning_task/lightning_task.py

@@ -93,7 +93,7 @@ def fit(
             self.trainer.fit(self.model, data)
         except Exception as e:
             del self.trainer
-            torch.cuda.empty_cache()  # type: ignore
+            torch.cuda.empty_cache()
             raise e
 
     @abc.abstractmethod