return types in quotes

Author: aditya0by0

chebifier/prediction_models/electra_predictor.py

@@ -45,7 +45,7 @@ def __init__(self, model_name: str, ckpt_path: str, **kwargs):
         super().__init__(model_name, ckpt_path, reader_cls=ChemDataReader, **kwargs)
         print(f"Initialised Electra model {self.model_name} (device: {self.device})")
 
-    def init_model(self, ckpt_path: str, **kwargs) -> Electra:
+    def init_model(self, ckpt_path: str, **kwargs) -> "Electra":
         from chebai.models.electra import Electra
 
         model = Electra.load_from_checkpoint(

chebifier/prediction_models/gnn_predictor.py

@@ -33,7 +33,7 @@ def load_class(self, class_path: str):
         module = __import__(module_path, fromlist=[class_name])
         return getattr(module, class_name)
 
-    def init_model(self, ckpt_path: str, **kwargs) -> ResGatedGraphConvNetGraphPred:
+    def init_model(self, ckpt_path: str, **kwargs) -> "ResGatedGraphConvNetGraphPred":
         import torch
         from chebai_graph.models.graph import ResGatedGraphConvNetGraphPred