bug for infoGAN when change data to celebA

[aflyingnoob]

infogan line 102 when I try to apply it on celebA with size (channel,height,width) as (3,180,180) as the code below, the data from dataloader size is (batch,channel,height,width), but when choose the [0] the data size become (3,180,180), so at the

self.D = discriminator(input_dim=data.shape[1]

the input_dim become 180 rather than 3,I remove the [0] to fix it

        data = self.data_loader.__iter__().__next__()[0]
        #  networks init
        self.G = generator(input_dim=self.z_dim, output_dim=data.shape[1], input_size=self.input_size, len_discrete_code=self.len_discrete_code, len_continuous_code=self.len_continuous_code)
        self.D = discriminator(input_dim=data.shape[1], output_dim=1, input_size=self.input_size, len_discrete_code=self.len_discrete_code, len_continuous_code=self.len_continuous_code)