How to define a static parameter in a model?

1

I’m trying to define a static tensor in a pytorch model. Initially I just created the variable as a tensor and then saved it in my model. However that leads to problem when I use, model.to(device), since the tensor does not get shifted to cuda devices when I do this.

I can make the parameter shift by defining it as a nn.parameter(), but I do not want to alter or train this parameter.

So what exactly is the best way to tell pytorch to ensure that this static tensor gets shifted between cuda and cpu, while not messing up anything with the optimizer?

2

module.register_buffer(“x”, tensor(…)), but for scalars a simple attribute (module.pi = 3.14) is sometimes better, both are auto-transfered to gpu.

1 Like