KeyError: tensor(0, requires_grad=True)

Thank you for posting the solution here @RX1 !