Hi!
I am trying to train two models in stages:
model1 = Model1()
model2 = Model2()
optimizer1 = Adam(model1.parameters())
optimizer2 = Adam(model2.parameters())
for epoch in range(epochs):
     for stage in range(stages):
         model2.eval()
         model1.train()
         for step in range(steps1):
              ....
              pred = model1()
              ....
              loss1.backward() # depends on model1 only
              optimizer1.step()
              optimizer1.zero_grad()
        model2.train()
        for step in range(steps2):
              ....
              with torch.no_grad():
                  encoded = model1.encoder(inp)
              projection = model1.proj(encoded)
              pred = model2(projection)
              ....
              loss2.backward()
              optimizer1.step()
              optimizer2.step()
              optimizer1.zero_grad()
              optimizer2.zero_grad()

And I am getting an CUDA error: an illegal memory access was encountered (embedding_dense_backward_cuda at /pytorch/aten/src/ATen/native/cuda/Embedding.cu:267)
(model1 is a BERT)
Is this an incorrect way to do this? How should I approach updating the model’s weights in this setting?