Upload model

modeling_phylogpn.py

@@ -242,11 +242,10 @@ class PhyloGPNModel(PreTrainedModel):
 
     def forward(self, input_ids: torch.Tensor):
         output_tensor = self._model(input_ids)
-        output_array = output_tensor.numpy(force=True)
 
         results = {}
 
         for idx, key in enumerate("ACGT"):
-            results[key] = output_array[:, :, idx]
+            results[key] = output_tensor[:, :, idx]
 
         return results