A main constraint of the current implementation is that each channel-group in a batch must have the same number of masked elements. I don't think this is a problem for us here, but its a note for the future.
@ivanzvonkov :
Also we currently don't include the latlons in the stacked elements. I think I originally left them out to prevent overfitting to locations, but we can run experiments with them both out and in.
I added einops to the environment for testing (its very lightweight), but can remove it (and skip the tests) if you'd prefer.
A main constraint of the current implementation is that each channel-group in a batch must have the same number of masked elements. I don't think this is a problem for us here, but its a note for the future.
@ivanzvonkov :
Also we currently don't include the latlons in the stacked elements. I think I originally left them out to prevent overfitting to locations, but we can run experiments with them both out and in.
I added einops to the environment for testing (its very lightweight), but can remove it (and skip the tests) if you'd prefer.