fixes bias shape

tfjgeorge · Feb 3, 2024 · 1d61aed · 1d61aed
1 parent 4ce6305
commit 1d61aed
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/nngeometry/generator/jacobian/grads.py b/nngeometry/generator/jacobian/grads.py
@@ -57,6 +57,7 @@ def Jv(cls, buffer, mod, layer, x, gy, v, v_bias):
  cls.flat_grad(buffer_flat, mod, layer, x, gy)
  v = v.view(-1)
  if v_bias is not None:
+ v_bias = v_bias.view(-1)
  v = torch.cat((v, v_bias))
  buffer.add_(torch.mv(buffer_flat, v))