Cleanup.

ludwig-ai · May 23, 2024 · 9b91d6b · 9b91d6b
1 parent a8e5376
commit 9b91d6b
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/ludwig/encoders/image/base.py b/ludwig/encoders/image/base.py
@@ -382,16 +382,17 @@ def __init__(
             raise ValueError("img_height and img_width should be identical.")
         self._input_shape = (in_channels, img_height, img_width)
 
+        config_dict: dict
         if use_pretrained and not saved_weights_in_checkpoint:
+            config_dict = {
+                "pretrained_model_name_or_path": pretrained_model,
+            }
             if output_attentions:
-                transformer = ViTModel.from_pretrained(
-                    pretrained_model_name_or_path=pretrained_model,
-                    attn_implementation="eager",
-                )
-            else:
-                transformer = ViTModel.from_pretrained(pretrained_model_name_or_path=pretrained_model)
+                config_dict["attn_implementation"] = "eager"
+
+            transformer = ViTModel.from_pretrained(**config_dict)
         else:
-            config_dict: dict = {
+            config_dict = {
                 "image_size": img_height,
                 "num_channels": in_channels,
                 "patch_size": patch_size,