Add no-grad block

openfold/model/primitives.py

@@ -149,6 +149,7 @@ class Linear(nn.Linear):
             with torch.no_grad():
                 self.bias.fill_(0)
 
+        with torch.no_grad():
             if init_fn is not None:
                 init_fn(self.weight, self.bias)
             else:
@@ -161,7 +162,6 @@ class Linear(nn.Linear):
                 elif init == "gating":
                     gating_init_(self.weight)
                     if bias:
-                    with torch.no_grad():
                         self.bias.fill_(1.0)
                 elif init == "normal":
                     normal_init_(self.weight)