Thanks, it seems to require an incredible amount of memory. My GPU has 8 GB of memory and prior to calling hessian I have only allocated 121 kb according to torch.cuda.memory_allocated(0).
I am attempting to calculate a 1x1 Hessian for 360 x 20 batches, so my input is of the form [320 x 20 x 1]
When I call hessian with vectorize=False, it creates the [320 x 20 x 1 x 320 x 20 x 1] as expected. If I try to call with vectorize=True, I get an out of memory error when it takes more than 5 GB.
Is this expected behavior? I assume a 320 x 20 x 1 matrix is far smaller than most models. Am I correct to assume the output will be 320 x 20 x 1 x 1? The documentation doesn’t really explain what the vectorize flag expected behavior is except that it uses the vmap
Thank you so much for you help. If its unclear why this requires so much memory I’ll try to figure out other options
EDIT: If I run a for loop over the 320 x 20 batches to calculate each 1 x 1 Hessian with create_graph=True, it takes about 2.8 GB