Explore areas for optimization with torch.vmap

[kafitzgerald]

https://docs.pytorch.org/docs/stable/generated/torch.vmap.html#torch-vmap

Potential areas:

• interpolation
• other operations that can be vectorized