How to use the input of first iter to init variable in module?

Question

I want to use something related to input to init the property in my module, before first iter it is init to zeros in init()

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

        self.test_var = torch.zeros((1,3))

    def forward(self, imgs):
  
        print(self.test_var)
        if torch.equal( self.test_var, torch.zeros(self.test_var.shape)):
            print('zeros. ')
            var = torch.tensor([1,1,1])
            self.test_var = var + self.test_var

It should change the value of test_var to [1,1,1] after 1st iter, but I only find the output is enter image description here It means the change of test_var is invalid, and it is set to zero every iter. I'm confused. I guess maybe something wrong with DataParallel, because when I set CUDA_VISIEBLE_DEVICES=0, the question can be fixed.

Answer 1

What piece of code do you use to call the model.forward() function?
Could it be that you reinitialize the model after each epoch?
Also, your imgs param is not used in the code you provided.

When I initialize the model and call model.forward() like this...

model = Model()
epochs = 4
for epoch in range(epochs):
  print(f"Epoch {epoch+1}:")
  model.forward(None)
  print("--")

...I get the following output that you also expect to come: In the first epoch, the ones are added because both tensors contain only zeros. In the following they are not added because the zero values are not equal to the one values.

Epoch 1:
tensor([[0., 0., 0.]])
zeros. 
--
Epoch 2:
tensor([[1., 1., 1.]])
--
Epoch 3:
tensor([[1., 1., 1.]])
--
Epoch 4:
tensor([[1., 1., 1.]])
--