I don’t get the error message using resnet50 as the pretrained_net.
However, you are currently not registering stage2 and stage3 properly, since you are using a Python list. Use nn.ModuleList instead. Otherwise, model.parameters() won’t return these parameters and model.to() won’t push the parameters to the specified device.