add license

lucidtronix · lucidtronix · commit 91dac6ed09dc · 2025-05-09T12:54:28.000-04:00
diff --git a/ml4h/models/train_diffusion.py b/ml4h/models/train_diffusion.py
@@ -131,7 +131,7 @@ def regress_on_batch(diffuser, regressor, controls, tm_out, batch_size):
         diffusion_steps=50,
     )
     control_predictions = regressor.predict(generated_images, verbose=0)
-    return control_predictions[:, 0] if tm_out.is_continuous() else control_predictions
+    return control_predictions[tm_out.output_name()][:, 0] if tm_out.is_continuous() else control_predictions
 
 
 def regress_on_controlled_generations(diffuser, regressor, tm_out, batches, batch_size, mean, std, prefix):

Original file line number	Diff line number	Diff line change
`@@ -131,7 +131,7 @@ def regress_on_batch(diffuser, regressor, controls, tm_out, batch_size):`
`131`	`131`	`diffusion_steps=50,`
`132`	`132`	`)`
`133`	`133`	`control_predictions = regressor.predict(generated_images, verbose=0)`
`134`		`- return control_predictions[:, 0] if tm_out.is_continuous() else control_predictions`
	`134`	`+ return control_predictions[tm_out.output_name()][:, 0] if tm_out.is_continuous() else control_predictions`
`135`	`135`
`136`	`136`
`137`	`137`	`def regress_on_controlled_generations(diffuser, regressor, tm_out, batches, batch_size, mean, std, prefix):`