WGAN on grayscale (1 channel) images

I’m trying to run WGAN based on lesson7-wgan jupyter notebook on grayscale (1 channel) images [MTI data]
To adjust the run from 3 to 1 channels I’ve modified the following:
from:
.normalize(stats = [torch.tensor([0.5,0.5,0.5]), torch.tensor([0.5,0.5,0.5])], do_x=False, do_y=True))
to:
.normalize(stats = [torch.tensor([0.5]), torch.tensor([0.5])], do_x=False, do_y=True))
from:
generator = basic_generator(in_size=64, n_channels=3, n_extra_layers=1)
critic = basic_critic (in_size=64, n_channels=3, n_extra_layers=1)
to:
generator = basic_generator(in_size=64, n_channels=1, n_extra_layers=1)
critic = basic_critic (in_size=64, n_channels=1, n_extra_layers=1)

And I still get the following error message:


What could be the problem?

Thanks
Moran