Added bias for conv: diagonal case

Author: IlMioFrizzantinoAmabile

stochman/nnj.py

@@ -711,6 +711,11 @@ def _jacobian_wrt_weight_sandwich_diag_to_diag(self, x: Tensor, val: Tensor, tmp
             output_tmp_single_batch = output_tmp_single_batch.reshape(c2 * c1 * kernel_h * kernel_w)
             output_tmp[i, :] = output_tmp_single_batch
 
+        if self.bias is not None:
+            bias_term = tmp_diag.reshape(b, c2, h2*w2)
+            bias_term = torch.sum(bias_term, 2)
+            output_tmp = torch.cat([output_tmp, bias_term], dim=1)
+
         return output_tmp