I am trying to train two tasks one after another in multitask learning. I have two CNNs for two tasks. While I am computing task 1 layers I also need to use CNN layers output from the other task. I am getting retain_graph = True error as I initialized all layers output as tensors! I don’t know how to fix it!
# Codes for https://www.aclweb.org/anthology/N18-2114/ / Gated Multi-Task Network for Text Classification
class GCNN_MTL(nn.Module):
    def __init__(self, lins, opt):
        super(GCNN_MTL, self).__init__()
        self.opt = opt
        self.tasks = 2
        self.cnns = clones(CNN_Part(opt.EMB_DIM, opt.NUM_HIDDEN, opt), self.tasks)
        self.linears = lins
        self.sig_gate = Sigma_Gate()
        self.conv_out = [[torch.zeros(self.opt.NUM_FILTER, self.opt.LENGTH - 1).cuda(),
                         torch.zeros(self.opt.NUM_FILTER, self.opt.LENGTH - 2).cuda(),
                         torch.zeros(self.opt.NUM_FILTER, self.opt.LENGTH - 3).cuda()]] * self.tasks
        self.pool_out = [[torch.zeros(self.opt.NUM_FILTER).cuda(), torch.zeros(self.opt.NUM_FILTER).cuda(),
                         torch.zeros(self.opt.NUM_FILTER).cuda()]] * self.tasks
        self.fc_out = [torch.zeros(self.opt.MID_DIM).cuda()] * self.tasks

    def forward(self, w_emb, num_task):
        self.conv_out[num_task] = self.cnns[num_task](w_emb)  # [B, Co, L] * K
        g_conv_out = [self.sig_gate(x, y) for x, y in zip(self.conv_out[num_task], self.conv_out[1 - num_task])]
        self.pool_out[num_task] = [F.max_pool1d(x, x.size(2)).squeeze(2) for x in g_conv_out]  # [B, Co] * K
        g_pool_out = [self.sig_gate(x, y) for x, y in zip(self.pool_out[num_task], self.pool_out[1 - num_task])]
        g_pool_out = torch.cat(g_pool_out, dim=1)  # [B, Co * K]

        self.fc_out[num_task] = self.linears[num_task](g_pool_out)
        gated_fc_out = self.sig_gate(self.fc_out[num_task], self.fc_out[1 - num_task])
        return gated_fc_out