modified: src/flux/model.py

src/flux/model.py

@@ -99,7 +99,8 @@ class Flux(nn.Module):
         bias_device = self.img_in.bias.device if self.img_in.bias is not None else "N/A (None)"
         #print(f"self.img_in.weight device: {weight_device}")
         #print(f"self.img_in.bias device: {bias_device}")
-
+        print("Model img_in weight sample:", self.img_in.weight[0, :10]) # Print first 10 elements of first row
+        print("Model img_in bias sample:", self.img_in.bias[:10]) # Print first 10 elements o
         # --- FIX: Explicitly move img to the device of img_in's weight if they differ ---
         # This is the core fix if the mismatch is here
         if img.device != weight_device: