Mirrors/the-algorithm-ml
2
0
Fork 0
mirror of https://github.com/twitter/the-algorithm-ml.git synced 2025-01-10 06:49:21 +01:00
Code Issues Packages Projects Releases Wiki Activity

Fix typo in log_weights.py comment

Browse Source
This commit is contained in:
Brian Jordan 2023-03-31 15:08:54 -04:00 committed by GitHub
parent 78c3235eee
commit 974d6458af
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
common/log_weights.py
View File

@ -74,7 +74,7 @@ def log_ebc_norms(
: min(sample_size, emb_weight_tensor.shape[0]) : min(sample_size, emb_weight_tensor.shape[0])
] ]
# WARNING: .cpu() transfer executes malloc that may be the cause of memory leaks # WARNING: .cpu() transfer executes malloc that may be the cause of memory leaks
# Change sample_size if the you observe frequent OOM errors or remove weight logging. # Change sample_size if you observe frequent OOM errors or remove weight logging.
norms = emb_weight_tensor[sample_mask].cpu().norm(dim=1).to(torch.float32) norms = emb_weight_tensor[sample_mask].cpu().norm(dim=1).to(torch.float32)
logging.info(f"Norm shape before reduction: {norms.shape}", rank=-1) logging.info(f"Norm shape before reduction: {norms.shape}", rank=-1)
norms = norms.mean().to(torch.device(f"cuda:{dist.get_rank()}")) norms = norms.mean().to(torch.device(f"cuda:{dist.get_rank()}"))