Update chatNT.py

chatNT.py

@@ -590,7 +590,6 @@ class TorchMultiOmicsModel(PreTrainedModel):
     def __init__(self, config: ChatNTConfig) -> None:
         print("(debug) Entering in class")
         if isinstance(config, dict):
-            print("(debug) going in if condition")
             # If config is a dictionary instead of ChatNTConfig (which can happen
             # depending how the config was saved), we convert it to the config
             config["gpt_config"]["rope_config"] = RotaryEmbeddingConfig(
@@ -598,14 +597,24 @@ class TorchMultiOmicsModel(PreTrainedModel):
             )
             config["gpt_config"] = GptConfig(**config["gpt_config"])
             config["esm_config"] = ESMTransformerConfig(**config["esm_config"])
-            print("(debug) Type esm_config : ", type(config["esm_config"]))
-            print("(debug) esm_config : ", config["esm_config"])
             config["perceiver_resampler_config"] = PerceiverResamplerConfig(
                 **config["perceiver_resampler_config"]
             )
             config = ChatNTConfig(**config)  # type: ignore
-            print("(debug) Type config : ", type(config))
 
+        else:
+            if isinstance(config.gpt_config, dict):
+                config.gpt_config["rope_config"] = RotaryEmbeddingConfig(
+                    **config.gpt_config.["rope_config"]
+                )
+                config.gpt_config = GptConfig(**config.gpt_config)
+
+            if isinstance(config.esm_config, dict):
+                config.esm_config = ESMTransformerConfig(**config.esm_config)
+
+            if isinstance(config.perceiver_resampler_config, dict):
+                config.esm_config = PerceiverResamplerConfig(**config.perceiver_resampler_config)
+                
         print("(debug) config : ", config)
         print("(debug) config type : ", type(config))
         print("(debug) gpt config : ", config.gpt_config)