What is the best practice to save and load a BlockMask object?

[complexfilter]

Hi, I encountered a Key Error 'm' bug after

1. saving a BlockMask object to a .pkl file
2. loading the object from the pkl file
3. migrating all the tensors in the object from GPU 0 to GPU 1.
4. using the object on GPU 1 and q k v on GPU 1 to do flexattention.

So, I was wondering what is the best practice to save a blockmask object to disc. And how to load this disk file to multiple GPUs.

[drisspg]

Out of curiosity why do you want to create the BlockMask object on 1 GPU and share as opposed to creating it on N times for each GPU?