AlphaBetaGamma96

@MaxH1996 Neither am I, although I did have a quick look ( [source here](https://pytorch.org/docs/stable/_modules/torch/nn/modules/batchnorm.html#BatchNorm2d) ) at it and there are indeed some in-place operations (I'm not sure if that's need...

@zou3519 Do you think it might be possible to use `tree_map` to update parameters? Instead of doing, ``` optim.zero_grad() loss.backward() optim.step() ```` Could you do something like this? ``` grad,...

@zou3519, Let's say I have some network that represents an R^N -> R^1 function, and I want to precondition the gradient of the loss by its Fisher information matrix (FIM)....

Hi @JoaoLages, if you want to compute multiple losses from a single sample you _should_ be able to do this via forward-mode AD. I've written a brief example below, ```...

> In my case the `model` is very heavy and I can only call `model.forward` with a single sample. I would like to first compute `Y = model.forward(X)` and then...

Have you tried using jacrev or jacfwd? Something like, ``` fnet, params = make_functional(text2imagemodel) #or the buffers version if your model uses buffers gradient = vmap(jacrev(fnet, argnums=0), in_dims=(None, 0))(params, X)...

> > Have you tried using jacrev or jacfwd? > > Something like, > > ``` > > fnet, params = make_functional(text2imagemodel) #or the buffers version if your model uses...

Perhaps it might be better to upgrade to the latest version? Also, doesn't `chunk_vmap` only work for chunking over a batch of samples? So, if you're working on a single-sample...

> > Perhaps it might be better to upgrade to the latest version? > > I'm in the latest PyPi version :) > > > Also, doesn't `chunk_vmap` only work...

> text = "a photo of an astronaut riding a horse on mars" > X = get_input_embeddings(text) # output [1, embedding_dim] > > Y = text2imagemodel(X) # shape [1, width,...