add scatter param

6a41c3e7 · zhuww · 16401cf5 · 6a41c3e7
Commit 6a41c3e7 authored Dec 08, 2022 by zhuww
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

fastfold/distributed/comm.py fastfold/distributed/comm.py +2 -2

No files found.
--- a/fastfold/distributed/comm.py
+++ b/fastfold/distributed/comm.py
@@ -136,11 +136,11 @@ class Copy(torch.autograd.Function):
        return _reduce(grad_output)
-def scatter(input: Tensor, dim: int = -1) -> Tensor:
+def scatter(input: Tensor, dim: int = -1, drop_unused=False) -> Tensor:
    if torch.is_grad_enabled() and input.requires_grad:
        input = Scatter.apply(input, dim)
    else:
-        input = _split(input, dim=dim)
+        input = _split(input, dim=dim, drop_unused=drop_unused)
    return input