MNIST dataset not supported correctly

[dimk1]

I created a new mnist.yaml and placed it under /config/attack/prototype/mnist.yaml. Its content is this:

num_workers: 4
pin_memory: True
non_blocking: True
prefetch: False
amp: False
device: cuda:0
client_optimizer: sgd
dataset: mnist
dataset_path: ./data
frequency_save: 0
batch_size: 128
lr: 0.01
lr_scheduler: CosineAnnealingLR
model: resnet
random_seed: 0
sgd_momentum: 0.9
wd: 0.0005
epochs: 10

I then run the following attack:

python ./attack/badnet.py --yaml_path ./config/attack/prototype/mnist.yaml --patch_mask_path ./resource/badnet/trigger_image.png --save_folder_name badnet_0_1

and I get the following error:

DEBUG:root:dataset_and_transform_generate done
DEBUG:root:Not DatasetFolder or ImageFolder, so iter through
DEBUG:root:Not DatasetFolder or ImageFolder, so iter through
DEBUG:root:Reminder: plz note that if p_num or pratio exceed the number of possible candidate samples
 then only maximum number of samples will be applied
DEBUG:root:Reminder: priority p_num > pratio, and choosing fix number of sample is prefered if possible 
INFO:root:poison num:6000.0,real pratio:0.1
2025-05-16:16:25:46 [INFO    ] [backdoor_generate_poison_index.py:61] poison num:6000.0,real pratio:0.1
DEBUG:root:poison train idx is saved
INFO:root:save file format is .png
2025-05-16:16:25:46 [INFO    ] [bd_dataset_v2.py:133] save file format is .png
prepro_backdoor:   0%|                                                                                                                                                                                                                                | 17/60000 [00:00<00:01, 59271.13it/s]
Traceback (most recent call last):
  File "./attack/badnet.py", line 275, in <module>
    attack.stage1_non_training_data_prepare()
  File "./attack/badnet.py", line 147, in stage1_non_training_data_prepare
    bd_train_dataset = prepro_cls_DatasetBD_v2(
  File "./utils/bd_dataset_v2.py", line 236, in __init__
    self.prepro_backdoor()
  File "./utils/bd_dataset_v2.py", line 247, in prepro_backdoor
    img = self.bd_image_pre_transform(img, target=label, image_serial_id=selected_index)
  File "./utils/aggregate_block/bd_attack_generate.py", line 29, in __call__
    img = transform(img, *args, **kwargs)
  File "./utils/bd_img_transform/patch.py", line 44, in __call__
    return self.add_trigger(img)
  File "./utils/bd_img_transform/patch.py", line 47, in add_trigger
    return img * (self.trigger_array == 0) + self.trigger_array * (self.trigger_array > 0)
**ValueError: operands could not be broadcast together with shapes (28,28) (28,28,3)** 

I presume that the code does not create grayscale triggers correctly when MNIST dataset is selected and tries to apply triggers with 3 channels (28,28,3) on grayscale images (28,28)