Changed docstring of readout return type: float -> torch.Tensor

choderalab · Jan 18, 2024 · 952abbd · 952abbd
1 parent 48ffe2e
commit 952abbd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mtenn/readout.py b/mtenn/readout.py
@@ -70,7 +70,7 @@ def forward(self, delta_g):
 
         Returns
         -------
-        float
+        torch.Tensor
             Calculated pIC50 value.
         """
         pic50 = -delta_g / torch.log(torch.tensor(10, dtype=delta_g.dtype))
@@ -116,7 +116,7 @@ def forward(self, delta_g):
 
         Returns
         -------
-        float
+        torch.Tensor
             Calculated Ki value.
         """
         ki = torch.exp(delta_g)