torch meshgrid behavior changing

[billbrod]

With torch version 2.1, calling torch.meshgrid raises the following warning: "torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument". See meshgrid docs for more info.

we only call torch.meshgrid in a handful of places, so should be easy enough to fix.