def generator_step()

in torchrecipes/vision/image_generation/module/infogan.py [0:0]

• 21 lines of code
• 2 McCabe index (conditional complexity)

    def generator_step(self) -> TrainOutput:
        # Sample noise and labels as generator input
        z = torch.tensor(
            np.random.normal(0, 1, (self.cur_batch_size, self.latent_dim)),
            dtype=torch.float,
            device=self.device,
        )

        label_input = to_categorical(
            np.random.randint(0, self.n_classes, self.cur_batch_size),
            num_columns=self.n_classes,
            device=self.device,
        )
        code_input = torch.tensor(
            np.random.uniform(-1, 1, (self.cur_batch_size, self.code_dim)),
            dtype=torch.float,
            device=self.device,
        )

        # Generate a batch of images
        self.gen_imgs = self.generator(z, label_input, code_input)

        # Loss measures generator's ability to fool the discriminator
        validity, _, _ = self.discriminator(self.gen_imgs)
        g_loss = adversarial_loss(validity, self.valid)
        self.log("g_loss", g_loss, prog_bar=True)
        return g_loss