Skip to content

Commit

Permalink
Reduce code and rely on device checking in tensor.to
Browse files Browse the repository at this point in the history
Co-authored-by: Mayank Mittal <[email protected]>
Signed-off-by: jtigue-bdai <[email protected]>
  • Loading branch information
jtigue-bdai and Mayankm96 authored Oct 11, 2024
1 parent 14bde30 commit 5b99208
Showing 1 changed file with 1 addition and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,7 @@ def constant_noise(data: torch.Tensor, cfg: noise_cfg.ConstantNoiseCfg) -> torch

# fix tensor device for bias on first call and update config parameters
if isinstance(cfg.bias, torch.Tensor):
if cfg.bias.device is not data.device:
cfg.bias = cfg.bias.to(
data.device,
)
cfg.bias = cfg.bias.to(device=data.device)

if cfg.operation == "add":
return data + cfg.bias
Expand Down

0 comments on commit 5b99208

Please sign in to comment.