Use enumerate to clean up ANIModel (#358)

b9422498 · Gao, Xiang · Farhad Ramezanghorbani · 669f70a1 · b9422498
Commit b9422498 authored Nov 05, 2019 by Gao, Xiang Committed by Farhad Ramezanghorbani Nov 05, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 3 deletions

torchani/nn.py torchani/nn.py +2 -3

No files found.
--- a/torchani/nn.py
+++ b/torchani/nn.py
@@ -41,10 +41,9 @@ class ANIModel(torch.nn.Module):

        output = torch.full(species_.shape, self.padding_fill,
                            dtype=aev.dtype, device=species.device)
-        i = 0
-        for m in self.module_list:
+
+        for i, m in enumerate(self.module_list):
            mask = (species_ == i)
-            i += 1
            midx = mask.nonzero().flatten()
            if midx.shape[0] > 0:
                input_ = aev.index_select(0, midx)