Teardown trying to copy "meta" tensors about pytorch-lightning HOT 2 CLOSED

Hey @kvndhrty
I think a pretty easy way to work around this is to not register your meta-template model as a submodule. You can easily do that by packing it into a list:

def __init__(self):
    super().__init__()
    with torch.device("meta"):
        self._template_model = [TemplateModel()]
        
    # then access it like so in your other code: 
    self._template_model[0]
    
    # ... or write a getter to return you the template model without indexing

I think that the assumption Lightning makes about your model not being on the meta device after training is a reasonable one. Even so before training, since eventually Lightning moves the model to GPU before training. I think it would become quite complex if we had to add logic to ignore such submodules on the meta-device. More so, it would be error-prone, because meta-device initialization is needed for large model training.
So I would like to suggest we don't treat this as a bug.

One other thing you could do is ask yourself whether it is even necessary to have your template model as an attribute at all. Since the creation on meta-device is basically free, you could also just do that on-the-fly whenever you need that. Get the properties you need and store them somewhere. Then you don't need to keep that template model around.

from pytorch-lightning.

@awaelchli I think that is entirely reasonable, I'll pack my module into a list for now. The small re-factor required to init the meta module each time isn't something I'll do this week, but maybe in the near future.

Thank you for the quick response!

from pytorch-lightning.