code: def set_input_shape_ngpu(self, new_input_shape):
      
      assert self.device_name, "Device name has not been set."
  
      device_name = self.device_name
      if self.input_shape is None:
        # First time setting the input shape
        self.input_shape = [None] + [int(d) for d in list(new_input_shape)]
  
      if device_name in self.params_device:
        # There is a copy of weights on this device
        self.__dict__.update(self.params_device[device_name])
        return
  
      # Stop recursion
      self.params_device[device_name] = {}
  
      # Initialize weights on this device
      with tf.device(device_name):
        self.set_input_shape(self.input_shape)
        keys_after = self.__dict__.keys()
        if self.params_names is None:
          # Prevent overriding training
          self.params_names = [k for k in keys_after if isinstance(
              self.__dict__[k], tf.Variable)]
        params = {k: self.__dict__[k] for k in self.params_names}
        self.params_device[device_name] = params
docs: Create and initialize layer parameters on the device previously set
      in self.device_name.
  
      :param new_input_shape: a list or tuple for the shape of the input.