Take advantage of torch.where broadcasting (#447)

torchani/aev.py

@@ -94,7 +94,7 @@ def compute_shifts(cell: Tensor, pbc: Tensor, cutoff: float) -> Tensor:
     reciprocal_cell = cell.inverse().t()
     inv_distances = reciprocal_cell.norm(2, -1)
     num_repeats = torch.ceil(cutoff * inv_distances).to(torch.long)
-    num_repeats = torch.where(pbc, num_repeats, torch.zeros_like(num_repeats))
+    num_repeats = torch.where(pbc, num_repeats, num_repeats.new_zeros(()))
     r1 = torch.arange(1, num_repeats[0] + 1, device=cell.device)
     r2 = torch.arange(1, num_repeats[1] + 1, device=cell.device)
     r3 = torch.arange(1, num_repeats[2] + 1, device=cell.device)