Update chatNT.py

chatNT.py

@@ -720,8 +720,11 @@ class TorchMultiOmicsModel(PreTrainedModel):
                     projected_bio_embeddings.append(proj)
                     for key in output.keys():
                         outs[f"{key}_{bio_seq_num}"] = output[key]
+                    outs[f"bio_embeddings_list_{bio_seq_num}"] = proj
+                    
                 
                 projected_bio_embeddings = torch.stack(projected_bio_embeddings, dim=1)
+                outs["projected_bio_embeddings"] = projected_bio_embeddings.clone()
 
         # decode        
         logits = self.biobrain_decoder(
@@ -730,13 +733,6 @@ class TorchMultiOmicsModel(PreTrainedModel):
         )
 
         outs["logits"] = logits
-        outs["projected_bio_embeddings"] = projected_bio_embeddings
-        
-        # Just for debugging
-        print("(debug) remember to remove bio_embeddings storage")
-        if projected_bio_embeddings is not None:
-            for i, embed in enumerate(bio_embeddings_list):
-                outs[f"bio_embeddings_list_{i}"] = embed
 
         return outs