def __init__(self, input, model, num_gpu=1):
        """
        Args:
            input (InputSource):
            model (GANModelDesc):
        """
        super(GANTrainer, self).__init__()
        assert isinstance(model, GANModelDesc), model

        if num_gpu > 1:
            input = StagingInput(input)

        # Setup input
        cbs = input.setup(model.get_input_signature())
        self.register_callback(cbs)

        if num_gpu <= 1:
            self._build_gan_trainer(input, model)
        else:
            self._build_multigpu_gan_trainer(input, model, num_gpu)
def __init__ ( self , input , model , num_gpu = 1 ) : super ( GANTrainer , self ) . __init__ ( ) assert isinstance ( model , GANModelDesc ) , model if num_gpu > 1 : input = StagingInput ( input ) cbs = input . setup ( model . get_input_signature ( ) ) self . register_callback ( cbs ) if num_gpu <= 1 : self . _build_gan_trainer ( input , model ) else : self . _build_multigpu_gan_trainer ( input , model , num_gpu )
Args:
            input (InputSource):
            model (GANModelDesc):