Redesign GriddedDataset with Pytorch idioms in mind

[iancze]

This is the second part of dataset redesign (first parts discussed in #126 and #162), focused specifically on the changes to GriddedDataset.

All of the same design challenges present in the UVDataset redesign (#162) also apply here. Additionally, we have the concern that we need to deal with masks.

One immediate problem with the current design of GriddedDataset is that it does not inherit from any Pytorch class. There is a workflow where it might make sense to attach the dataset as an instance attribute of some class (inherited from nn.module), and then use the .to(device) syntax to send everything, including the dataset, to the GPU. We tried to duck type this behavior with the GriddedDataset.to method, but I don't think this entirely meshes with this workflow. We would like to have seamless transfer/to from CPU/GPU using Pytorch idioms.

And, thinking further down the line, we would like this dataset object to play well with cross-validation.

TODO: flesh this issue out with more details on constraints/desired features.