Hi!
I have some parameters that have to be in-place pruned in each iteration so that the shape is changed, but that the model and the optimizer keep references to them. Though the forward pass succeeds after the pruning, the backward pass fails because the gradient function uses the old shape of the gradient.
This code below demonstrates the problem, resulting in the error message
RuntimeError: Function SumBackward0 returned an invalid gradient at index 0 - got [16, 64, 64] but expected shape compatible with [32, 64, 64].
I have tried to work around this by trying to delete the previous graph by removing the output from the scope (by setting REQUEST_GRAPH_DELETION = True in the code), but it is not guaranteed to always work.
I have tried many things to make PyTorch correctly update the metadata, but nothing worked reliably. Should I report it on GitHub?
import torch, time

REQUEST_GRAPH_DELETION = False  # if False, the error always occurs, if True, only sometimes
TEST_REPEAT_COUNT = 0  # 200


def run_loop():
    x = torch.ones((32, 64, 64), requires_grad=True)
    out = None

    while x.shape[0] > 0:
        yield x.shape  # for printing

        if REQUEST_GRAPH_DELETION:
            del out  # time.sleep(0.0001) helps, but shoudn't be relied on

        out = x.sum()
        assert out.grad_fn(out).shape == x.shape  # correct gradient shape
        out.backward()  # incorrect gradient shape because some metadata is not updated

        with torch.no_grad():  # change the shape
            N = x.shape[0] // 2

            grad = x.grad
            x.set_(x.data[:N])
            x.grad = grad[:N]


for shape in run_loop():
    print(shape)

if TEST_REPEAT_COUNT > 0:  # repeats
    error_count, succ_count = 0, 0
    for i in range(TEST_REPEAT_COUNT):
        try:
            shapes = []
            for shape in run_loop():
                shapes.append(shape[0])
            print("OK   ", shapes)
            succ_count += 1
        except:
            print("ERROR", shapes)
            error_count += 1
    print(f"{error_count} errors, {succ_count} successes")