Skip to content

Commit

Permalink
Fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
pablomlago committed Jan 11, 2025
1 parent c550796 commit 3888df4
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 2 deletions.
2 changes: 1 addition & 1 deletion src/brevitas/graph/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ def __init__(
rot_mat: torch.nn.Parameter,
rot_func: Callable,
axis: int,
K: Optional[int],
K: Optional[int] = None,
) -> None:
super().__init__()
self.rot_mat = rot_mat
Expand Down
2 changes: 1 addition & 1 deletion tests/brevitas/graph/test_equalization.py
Original file line number Diff line number Diff line change
Expand Up @@ -452,7 +452,7 @@ def test_apply_rotate(
generator = torch.Generator()
generator.manual_seed(SEED)
# Clone generator to make sure we can use the same rotation matrices
generator_clone = generator.clone_state()
generator_clone = copy.deepcopy(generator)

# Apply rotations on the model with unfused rotations
regions_unfused = list(
Expand Down

0 comments on commit 3888df4

Please sign in to comment.