Cheers - I took out all the .data calls and something changed - all the variables in EMM_NTM (memory, wr, ww) are now nan (which I’m guessing comes from trying to backprop, though I’m not sure). I registered a backward hook and it looks like a bunch of gradients/parameters for nn.Linear modules are getting initialized to nan which is weird.
Still worrisome is that the state doesn’t seem to change at all after 1-100 call(s) to optimizer.step() even with a large learning rate. The way I’m testing this is to save the state_dict at the very beginning of the training process and then compare it with the state_dict later on during training.
The core of my training loop looks like this:
ntm = NTM(num_inputs, num_hidden, num_inputs, batch, num_reads=1)

    try:
        ntm.load_state_dict(torch.load("models/copy_seqlen_{}.dat".format(seq_len)))
    except FileNotFoundError or AttributeError:
        pass

    ntm.train()

    state = ntm.state_dict()

    criterion = nn.MSELoss()
    optimizer = optim.RMSprop(ntm.parameters(), lr=5e-3, weight_decay=0.0005)

    max_seq_len = 20
    for length in range(10, max_seq_len):

        test = CopyTask(length, [num_inputs, 1], num_samples=2e4)

        data_loader = DataLoader(test, batch_size=batch, shuffle=True, num_workers=4)

        for epoch in range(5):

            for i, data in enumerate(data_loader, 0):
                inputs, labels = data
                inputs = Variable(inputs, requires_grad=True)
                labels = Variable(labels)

                optimizer.zero_grad()
                ntm.zero_grad()
                outputs = ntm(inputs)

                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()

                assert not (ntm.state_dict()['hid_to_out.bias'] == state['hid_to_out.bias'])[0]  # this just breaks it on the first loop

# do stuff with the outputs, plot, running loss, etc.