fix

fastvideo/attention/layer.py

@@ -48,10 +48,6 @@ class DistributedAttention(nn.Module):
                                   **extra_impl_args)
         # Register attn_impl as submodule if it has learnable parameters (e.g., SLA's proj_l)
         # This ensures its parameters are included in state_dict() for saving/loading
-        if isinstance(self.attn_impl, nn.Module):
-            self.add_module('attn_impl', self.attn_impl)
-        # Register attn_impl as submodule if it has learnable parameters (e.g., SLA's proj_l)
-        # This ensures its parameters are included in state_dict() for saving/loading
         if isinstance(self.attn_impl, nn.Module):
             self.add_module('attn_impl', self.attn_impl)
         self.num_heads = num_heads

fastvideo/layers/fp4linear.py

@@ -86,4 +86,4 @@ def fp4_linear_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso
     bias = self.bias if not self.skip_bias_add else None
     output = _LinearFWD4BWD16Fn.apply(x, self.weight, bias, "cutlass", 16, True)
     output_bias = self.bias if self.skip_bias_add else None
-    return output, output_bias
+    return output, output_bias