Fix bias dtype issue in mixed ops. (#11293)

Author: comfyanonymous

comfy/ops.py

@@ -504,10 +504,7 @@ def __init__(
 
                 self.in_features = in_features
                 self.out_features = out_features
-                if bias:
-                    self.bias = torch.nn.Parameter(torch.empty(out_features, **self.factory_kwargs))
-                else:
-                    self.register_parameter("bias", None)
+                self._has_bias = bias
 
                 self.tensor_class = None
                 self._full_precision_mm = MixedPrecisionOps._full_precision_mm
@@ -536,6 +533,10 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata,
                     self.weight = torch.nn.Parameter(weight.to(device=device, dtype=dtype), requires_grad=False)
                     if dtype != MixedPrecisionOps._compute_dtype:
                         self.comfy_cast_weights = True
+                    if self._has_bias:
+                        self.bias = torch.nn.Parameter(torch.empty(self.out_features, device=device, dtype=dtype))
+                    else:
+                        self.register_parameter("bias", None)
                 else:
                     self.quant_format = layer_conf.get("format", None)
                     if not self._full_precision_mm:
@@ -565,6 +566,11 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata,
                         requires_grad=False
                     )
 
+                    if self._has_bias:
+                        self.bias = torch.nn.Parameter(torch.empty(self.out_features, device=device, dtype=MixedPrecisionOps._compute_dtype))
+                    else:
+                        self.register_parameter("bias", None)
+
                     for param_name in qconfig["parameters"]:
                         param_key = f"{prefix}{param_name}"
                         _v = state_dict.pop(param_key, None)