In compute_fim_loss (see here), the code will only work if fim_reduce == "trace_max". @Rohan138 git blame at you, can you update this func as necessary to handle other fim_reduce values (if None, then what? throw exception?). Please add a docstring too.
Rohan138
commented
4 months ago
In
compute_fim_loss(see here), the code will only work iffim_reduce == "trace_max". @Rohan138 git blame at you, can you update this func as necessary to handle otherfim_reducevalues (ifNone, then what? throw exception?). Please add a docstring too.