Changed boolean statement to work with arrays

[Ubiquinone-dot]

Fixes the following error which I ran into: ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

The k1 array passed in this case is of shape (batch_size, n_atoms, 3)