Persistance of weight values between GPU and CPU

1

Hey!

I need a quick clarification to a simple question. Lets say I run the following code:

def train(net):
    net = net.cuda()
    for epoch in range(5):
        #train net and gather losses
    return losses

net = MyNet()
train(net)
torch.save(net.state_dict(), 'path/to/file')

Will the weights that are saved be the ones that are trained or the originally initialized weights?

How does weight values vary between the GPU and CPU version of the same neural-net object?

Thanks in advance!

1 Like
2

they will be the trained weights. (considering that “train net” consists of loss.backward() and optimizer.step() somewhere)

3

Thanks for the super quick reply! Pytorch is amazing, thanks for all the work you guys put in.