Simply reference weight handling for custom vs default blocks

Author: rwightman

src/open_clip/transformer.py

@@ -264,6 +264,9 @@ def __init__(
         ]))
         self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
 
+    def get_reference_weight(self):
+        return self.mlp.c_fc.weight
+
     def attention(
             self,
             q_x: torch.Tensor,
@@ -516,18 +519,10 @@ def __init__(
             ])
 
     def get_cast_dtype(self) -> torch.dtype:
-        # Handle both ResidualAttentionBlock and CustomResidualAttentionBlock
-        if hasattr(self.resblocks[0], 'get_reference_weight'):
-            # CustomResidualAttentionBlock has get_reference_weight method
-            weight = self.resblocks[0].get_reference_weight()
-            if hasattr(weight, 'int8_original_dtype'):
-                return weight.int8_original_dtype
-            return weight.dtype
-        else:
-            # ResidualAttentionBlock
-            if hasattr(self.resblocks[0].mlp.c_fc, 'int8_original_dtype'):
-                return self.resblocks[0].mlp.c_fc.int8_original_dtype
-            return self.resblocks[0].mlp.c_fc.weight.dtype
+        weight = self.resblocks[0].get_reference_weight()
+        if hasattr(weight, 'int8_original_dtype'):
+            return weight.int8_original_dtype
+        return weight.dtype
 
     def forward_intermediates(
             self,