fix

sayakpaul · sayakpaul · commit 1d1e7157ca65 · 2025-05-01T18:22:28.000+05:30
diff --git a/tests/models/transformers/test_models_transformer_hidream.py b/tests/models/transformers/test_models_transformer_hidream.py
@@ -82,6 +82,7 @@ def prepare_init_args_and_inputs_for_common(self):
             "axes_dims_rope": (4, 2, 2),
             "max_resolution": (32, 32),
             "llama_layers": (0, 1),
+            "force_inference_output": True,  # TODO: as we don't implement MoE loss in training tests.
         }
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict

Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,7 @@ def prepare_init_args_and_inputs_for_common(self):`
`82`	`82`	`"axes_dims_rope": (4, 2, 2),`
`83`	`83`	`"max_resolution": (32, 32),`
`84`	`84`	`"llama_layers": (0, 1),`
	`85`	`+ "force_inference_output": True, # TODO: as we don't implement MoE loss in training tests.`
`85`	`86`	`}`
`86`	`87`	`inputs_dict = self.dummy_input`
`87`	`88`	`return init_dict, inputs_dict`