Deploying Shiny app with custom Keras Models

I am trying to create a shiny app that uses a custom keras model.

Here is an example of a custom layer I have created:

Downsampling(keras$layers$Layer) %py_class% {

  initialize <- function(out_channels, first_stride = 1, kernelsize = c(4,4),
                         firstpadding = "valid") {
    super$initialize()
    self$out_channels <- out_channels
    self$first_stride <- first_stride
    self$kernelsize <- kernelsize
    self$first_padding = firstpadding



  }

  build <- function(input_shape) {
    self$conv_sequence <- keras_model_sequential() %>%
      layer_conv_2d(filters = self$out_channels, kernel_size = self$kernelsize ,
                    strides = self$first_stride, padding = self$first_padding,
                    kernel_regularizer=regularizer_l2(0.01),
                    use_bias = FALSE, trainable = TRUE) %>%
      layer_activation_leaky_relu(alpha = 0.2, trainable = TRUE) %>%
      layer_max_pooling_2d(pool_size = c(2,2)) %>%
      layer_batch_normalization(trainable = TRUE)

  }


  call <- function(self, inputs) {
    x <- self$conv_sequence(inputs)

    return(x)

  }

  get_config <- function() {
    list(out_channels = self$out_channels,
         first_stride = self$first_stride,
         kernelsize = self$kernelsize)
  }
}

layer_downsampling <- keras::create_layer_wrapper(Downsampling)

I also have a similar upsampling layer. I create the model using the keras_model function and save it

generator <- keras_model(gi, f)
save_model_tf(generator, paste0("DeconTool/GenModel/Model/Model"))

In source file, I can then load it. Note, the source file has the generator model written (so it has all the custom layers and how they connect etc.).

generator <- load_model_tf("GenModel/Model/Model", 
                             custom_objects = list("generator" = generator))

The code works fine on my local machine. I can upload it to shiny, but I get an error when I try to open the web app.

The log file gives an warning and then an error

Blockquote
2024-03-28 16:13:32.910382: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT

Blockquote
2024-03-28T16:13:34.500048+00:00 shinyapps[11607313]: Error in py_call_impl(callable, call_args$unnamed, call_args$named) :
2024-03-28T16:13:34.504492+00:00 shinyapps[11607313]: ValueError: For a build() method with more than one argument, all arguments should have a _shape suffix and match an argument from call(). E.g. build(self, foo_shape, bar_shape) For layer 'Downsampling', Received build() argument self, which does not end in _shape.

Any help would be greatly appreciated!