Formatting

packages/pidml/pidml/networks.py

@@ -162,8 +162,9 @@ class SimpleNet(GeneralNet):
             'probabilities' that each particle was of each possible particle
             type.
         """
+        n = self.n_detector
         outs = [
-            self.fcs[i](x[:, i * self.n_detector : (i + 1) * self.n_detector])
+            self.fcs[i](x[:, i * n : (i + 1) * n])
             for i in range(self.n_class)
         ]
         out = torch.cat(outs, dim=1)