I found that there is still possible to reduce the memory which is a little similar to gradient checkpoint. This paper 6 proposes inplace activated batchnorm, which only need to save the ouput for backward. While using imperative PyTorch implementation bn + relu need to save two tensor for bn and relu respectively, that method only need to save one tensor. Nevertheless, it demand the activation function to be invertible, which is not the case for relu…   But leaky relu supports it.