fix formatting

SkafteNicki · web-flow · commit f510a874c5bd · 2022-06-09T10:15:18.000+02:00
diff --git a/stochman/nnj.py b/stochman/nnj.py
@@ -665,7 +665,7 @@ def _jacobian_wrt_weight_sandwich_full_to_full(self, x: Tensor, val: Tensor, tmp
         )
 
     def _jacobian_wrt_weight_sandwich_full_to_diag(self, x: Tensor, val: Tensor, tmp: Tensor) -> Tensor:
-        ### TODO: Implement this in a smarter way
+        # TODO: Implement this in a smarter way
         return torch.diagonal(self._jacobian_wrt_weight_sandwich_full_to_full(x, val, tmp), dim1=1, dim2=2)
 
     def _jacobian_wrt_weight_sandwich_diag_to_full(self, x: Tensor, val: Tensor, tmp_diag: Tensor) -> Tensor:

Original file line number	Diff line number	Diff line change
`@@ -665,7 +665,7 @@ def _jacobian_wrt_weight_sandwich_full_to_full(self, x: Tensor, val: Tensor, tmp`
`665`	`665`	`)`
`666`	`666`
`667`	`667`	`def _jacobian_wrt_weight_sandwich_full_to_diag(self, x: Tensor, val: Tensor, tmp: Tensor) -> Tensor:`
`668`		`- ### TODO: Implement this in a smarter way`
	`668`	`+ # TODO: Implement this in a smarter way`
`669`	`669`	`return torch.diagonal(self._jacobian_wrt_weight_sandwich_full_to_full(x, val, tmp), dim1=1, dim2=2)`
`670`	`670`
`671`	`671`	`def _jacobian_wrt_weight_sandwich_diag_to_full(self, x: Tensor, val: Tensor, tmp_diag: Tensor) -> Tensor:`