Hello there,
I’m currently working on a GAN that has 2 paired RNNs as generator (rnn and output in the following code).
My problem is that the weights of the generator are not being updated.
The following code is used to perform a sigle forward step of the 2 RNNs to generate a single observation.
The main idea is that I have the tensor y_pred_long that is filled step after step with the predictions of the 2 RNNs.
Then this structure (y_pred_long) is passed to a decoding function that retrieves the desidered input for the discriminator.
I not beeing able to understand whether/where the computational graph breaks down, any suggestion is very appreciated!
If anything from the following code is not clear and additional information is needed please let me know! (first time posting here)

def gen_rnn_epoch(epoch, rnn, output, device):

    rnn.hidden = rnn.init_hidden(batch_size=1)  # this is the source of noise for the GAN generator

    # empty tensor that will be filled at each step with predictions from both RNNs
    y_pred_long = torch.zeros(1, 9, 8 * 5 + 4).to(device)
    # the first input of RNN_1 is a token
    x_step = torch.ones(1, 1, 8 * 5 + 4).to(device)

    # sequential generation of RNN_1
    for i in range(9):

        h, node_prediction = rnn(x_step)
        node_prediction = F.gumbel_softmax(node_prediction, dim=2, hard=True)

        # init RNN_2: context init
        hidden_null = torch.zeros(rnn.num_layers - 1, h.size(0), h.size(2)).to(device)
        output.hidden = torch.cat((h.permute(1, 0, 2), hidden_null), dim=0)

        # reset token and fill with newest prediction
        x_step = torch.zeros(1, 1, 8 * 5 + 4).to(device)
        x_step[:, :, :4] = node_prediction

	    # the first input of RNN_2 is a token
        output_x_step = torch.ones(1, 1, 5).to(device)
	
		# sequential generation of RNN_2
        edge_rnn_step = 0
        for j in range(min(8, i + 1)):
            output_y_pred_step = output(output_x_step)
            output_x_step = F.gumbel_softmax(output_y_pred_step, dim=2, hard=True)
            x_step[:, :, 4 * j + 4 + j: 4 * (j + 1) + 4 + (j + 1)] = output_x_step
            output.hidden = output.hidden.to(device)
            edge_rnn_step = j

        y_pred_long[:, i:i + 1, :] = x_step

        # stop criteria (EOS met)
        node_to_break, edges_to_break = torch.split(x_step, [4, 40], dim=2)
        edges_to_break_temp = torch.reshape(edges_to_break, (edges_to_break.shape[0], 8, 5))
        edges_to_break_uptillnow = edges_to_break_temp[0, :edge_rnn_step + 1, :]
        break_ = True
        for row in edges_to_break_uptillnow:
            if torch.argmax(row).item() != 0:
                break_ = False
            if edge_rnn_step == 0:
                break_ = False

        if break_:
            break

        rnn.hidden = rnn.hidden.to(device)

    x, edg_idx, edg_attr = decode_adj_generation(y_pred_long[0], device)

    data = Data(x=x, edge_index=edg_idx, edge_attr=edg_attr)

    return data

def decode_adj_generation(encoded_adj, device):

    max_prev_TEST = encoded_adj.shape[0] - 1  
    node_x, edges = torch.split(encoded_adj, [4, 40], dim=1)
    edges = torch.reshape(edges, (encoded_adj.shape[0], max_prev_TEST, 5))

    # drop all the rows for which we didn't have sampled a node
    idxs_to_keep = []
    for row in range(node_x.shape[0]):
        if (torch.sum(node_x[row]).item()) != 0:
            idxs_to_keep.append(row)
        else:
            break

    idxs_to_keep = torch.LongTensor(idxs_to_keep).to(device)
    node_x = node_x.index_select(0, idxs_to_keep)
    edges = edges.index_select(0, idxs_to_keep)

    # un-FLIP the edges
    for i in range(edges.shape[0]):
        edges[i, :i + 1, :] = torch.flip(edges[i, :i + 1, :], [0]) 

    # below for edge index creation, this step CAN be NON-differentiable
    first_dim_idx = []
    second_dim_idx = []

    for i in range(edges.shape[0]):
        for j in range(i):
            if torch.sum(edges[i - 1, j, :]) != 0:
                if np.nonzero(edges[i - 1, j, :])[0].item() != 0:
                    first_dim_idx.append(i)
                    second_dim_idx.append(j)

    first_dim_idx = torch.tensor(torch.LongTensor(first_dim_idx)).to(device)
    second_dim_idx = torch.tensor(torch.LongTensor(second_dim_idx)).to(device)
    row1 = torch.unsqueeze(torch.cat((first_dim_idx, second_dim_idx), 0), 0)
    row2 = torch.unsqueeze(torch.cat((second_dim_idx, first_dim_idx), 0), 0)
    edge_idx = torch.cat((row1, row2), 0)
    half_edge_idx = torch.cat((torch.unsqueeze(second_dim_idx, 0), torch.unsqueeze(first_dim_idx - 1, 0)), 0)

    # construction of edge_attr matrix, this has to be differentiable
    edge_attr = None
    step = 0
    for i in range(first_dim_idx.shape[0]):
        if step == 0:
            edge_attr = torch.unsqueeze((edges[half_edge_idx[1][i].item(), half_edge_idx[0][i].item(), 1:]), 0)
            step += 1
        else:
            edge_attr = torch.cat(
                (edge_attr, torch.unsqueeze(edges[half_edge_idx[1][i].item(), half_edge_idx[0][i].item(), 1:], 0)), dim=0)

    if edge_attr is not None:
        edge_attr = torch.cat((edge_attr, edge_attr), 0)
    else:
        edge_attr = torch.empty(0, 4)

    return node_x.to(torch.float32).to(device), edge_idx.to(torch.long).to(device), edge_attr.to(torch.float32).to(
            device)