CASTER: typing fix

AstraZeneca · Feb 3, 2022 · a9aaae8 · a9aaae8
1 parent 8205759
commit a9aaae8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chemicalx/models/caster.py b/chemicalx/models/caster.py
@@ -107,7 +107,7 @@ def dictionary_encoder(
         r = drug_pair_features_latent[:, None, :].matmul(dict_feat_closed_form.transpose(2, 1)).squeeze(1)
         return r
 
-    def forward(self, drug_pair_features: torch.FloatTensor) -> Tuple[torch.FloatTensor]:
+    def forward(self, drug_pair_features: torch.FloatTensor) -> Tuple[torch.FloatTensor, ...]:
         """Run a forward pass of the CASTER model.
 
         :param drug_pair_features: functional representation of each drug pair (see unpack method)