condition and supervise

broadinstitute · Jan 13, 2025 · d65c3e9 · d65c3e9
1 parent 75e86e6
commit d65c3e9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ml4h/models/train.py b/ml4h/models/train.py
@@ -375,7 +375,7 @@ def train_diffusion_control_model(args, supervised=False):
     if args.inspect_model:
         data, labels, paths = big_batch_from_minibatch_generator(generate_test, 1)
 
-        model.plot_reconstructions(data, prefix=f'{args.output_folder}/{args.id}/')
+        model.plot_reconstructions((data, labels), prefix=f'{args.output_folder}/{args.id}/')
         images = data[args.tensor_maps_in[0].input_name()]
         predictions_to_pngs(images, args.tensor_maps_in, args.tensor_maps_in, data, labels, paths, '{args.output_folder}/{args.id}/')
         interpolate_controlled_generations(model, args.tensor_maps_out, args.tensor_maps_out[0], args.batch_size,