diff --git a/stoked/forces.py b/stoked/forces.py index 36f6aef..9d7384d 100644 --- a/stoked/forces.py +++ b/stoked/forces.py @@ -17,7 +17,7 @@ def force(self): r_ij = self.position[:,np.newaxis] - self.position[np.newaxis] # N x N x 3 F_ij = self.force_func(r_ij) # N x N x 3 - np.einsum('iix->x', F_ij)[...] = 0 + np.einsum('iix->ix', F_ij)[...] = 0 F_i = np.sum(F_ij, axis=1) return F_i