Bug: Porting model another device (CPU/GPU)

[satabios]

• snntorch version: 0.7.0
• Python version: 3.10.13
• Operating System: Ubuntu x64

Description

When I load a pre-trained model and push it from GPU->CPU or the vice versa. Certain variables are still projecting in the original device and are not getting pushed to the destination device.

Specifically the "mem" variable under snntorch._neurons.leaky.Leaky

snn_pretrained_model_path = "snn_model.pth"
snn_model.load_state_dict(torch.load(snn_pretrained_model_path))  
snn_model.to("cpu") # or "gpu"

What I Did

As a workaround, I deliberately iterate over all the model to find such instances of leaky and push them to the destination device.

if isinstance(layer, nn.Sequential):
    for layer_id in range(len(original_dense_model)):
        layer = original_dense_model[layer_id]
        if isinstance((layer), snntorch._neurons.leaky.Leaky):
            layer.mem = layer.mem.to("cpu") # or "gpu" depending on the destination device
else:
    for internal_layer in model.modules():
        if isinstance((internal_layer), snntorch._neurons.leaky.Leaky):
            internal_layer.mem = internal_layer.mem.to("cpu")

[ahenkes1]

@satabios , have you tried the following steps?
Saving and loading models across devices in PyTorch

[satabios]

I followed Saving and loading models across devices in PyTorch.

The issue isn't during saving or reloading. However, it is with porting from one device to another, as mentioned above, Certain members of the model aren't getting transferred intrinsically.

[ahenkes1]

So the issue is there with newly created models? Without the saving/loading.

[satabios]

Either way, the bug persists. When the model is in memory and is prompted to transfer to a different device (say from GPU to CPU or vice versa) or when the model is loaded from a file. The porting causes the variables to be struck in the original device.