diff --git a/docs/anima_train_network.md b/docs/anima_train_network.md
index f97aa975..5d67ae36 100644
--- a/docs/anima_train_network.md
+++ b/docs/anima_train_network.md
@@ -652,4 +652,4 @@ The following metadata is saved in the LoRA model file:
 * `ss_sigmoid_scale`
 * `ss_discrete_flow_shift`
 
-</details>
+</details>
\ No newline at end of file
diff --git a/networks/lora_anima.py b/networks/lora_anima.py
index 224ef20c..9413e8c8 100644
--- a/networks/lora_anima.py
+++ b/networks/lora_anima.py
@@ -636,4 +636,4 @@ class LoRANetwork(torch.nn.Module):
             scalednorm = updown.norm() * ratio
             norms.append(scalednorm.item())
 
-        return keys_scaled, sum(norms) / len(norms), max(norms)
+        return keys_scaled, sum(norms) / len(norms), max(norms)
\ No newline at end of file
diff --git a/networks/lora_flux.py b/networks/lora_flux.py
index d74d0172..947733fe 100644
--- a/networks/lora_flux.py
+++ b/networks/lora_flux.py
@@ -141,10 +141,13 @@ class LoRAModule(torch.nn.Module):
             # rank dropout
             if self.rank_dropout is not None and self.training:
                 mask = torch.rand((lx.size(0), self.lora_dim), device=lx.device) > self.rank_dropout
-                if len(lx.size()) == 3:
-                    mask = mask.unsqueeze(1)  # for Text Encoder
-                elif len(lx.size()) == 4:
-                    mask = mask.unsqueeze(-1).unsqueeze(-1)  # for Conv2d
+                if isinstance(self.lora_down, torch.nn.Conv2d):
+                    # Conv2d: lora_dim is at dim 1 → [B, dim, 1, 1]
+                    mask = mask.unsqueeze(-1).unsqueeze(-1)
+                else:
+                    # Linear: lora_dim is at last dim → [B, 1, ..., 1, dim]
+                    for _ in range(len(lx.size()) - 2):
+                        mask = mask.unsqueeze(1)
                 lx = lx * mask
 
                 # scaling for rank dropout: treat as if the rank is changed
@@ -1445,4 +1448,4 @@ class LoRANetwork(torch.nn.Module):
             scalednorm = updown.norm() * ratio
             norms.append(scalednorm.item())
 
-        return keys_scaled, sum(norms) / len(norms), max(norms)
+        return keys_scaled, sum(norms) / len(norms), max(norms)
\ No newline at end of file