Skip to content

Commit

Permalink
fix (bias_corr/accelerate): Added comment
Browse files Browse the repository at this point in the history
  • Loading branch information
nickfraser authored Jul 8, 2024
1 parent 930296e commit b1d7d5a
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions src/brevitas/graph/calibrate.py
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,7 @@ def apply_correction(self, model):
if module.bias is not None:
module.bias.data += correction
elif self.skip_if_no_bias is False:
# If accelerate is enabled, bias will be on the same execution device as the weights, but won't be managed properly by accelerate
module.register_parameter(
'bias', nn.Parameter(correction).to(module.weight.device))
# Offload params again
Expand Down

0 comments on commit b1d7d5a

Please sign in to comment.