Propose to substitute the following line: https://github.com/cybertronai/autograd-hacks/blob/1fa974aa32a50aeadec449147acc5ae5a032f848/autograd_hacks.py#L165 with: `A = torch.nn.functional.unfold(A, layer.kernel_size, dilation=layer.dilation, padding=layer.padding, stride=layer.stride)`