@ptrblck I still don’t know how to implement the dimension reduction for my custom max operation as controlled by the dim parameter of torch.max but this reply made it clear that if the code runs without errors, then autograd handles the in-place operations.