add scatter param

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