Creating network layer by layer throws an error (add_ragged_softmax)

Description

Hi, guys!
I am building a network layer by layer, which contains add_ragged_softmax and it’s throwing an error:

[TRT] [E] 4: [raggedSoftMaxNode.cpp::computeOutputExtents::13] Error Code 4: Internal Error ((Unnamed Layer 40) [Ragged SoftMax]: Input tensor must have exactly 2 dimensions)*

on TensorRT8, but works without any problem on TensorRT7 (7.1) (need to build this engine on the machine with TensorRT8).

Here’s how I am creating the network and adding the layers which goes as an input to ragged softmax layer:

with trt.Builder(self.trt_logger) as builder, builder.create_network() as network, builder.create_builder_config() as builder_config:
.....

.....
add_bias2 = network.add_elementwise(matrix_multiplication2.get_output(0), bias2.get_output(0),
                                    trt.ElementWiseOperation.SUM) # shape of output from this operation is (30, 37) 
dimensions = network.add_input('dimensions', trt.int32, (30, 1))
softmax = network.add_ragged_softmax(add_bias2.get_output(0), dimensions)

builder.max_batch_size = self.max_batch_size
builder_config.max_workspace_size = 1 << 30
network.mark_output(softmax.get_output(0))
engine = builder.build_engine(network, builder_config)

Weights of the layer add_bias2 have builded correctly with small epsilon.

Environment

TensorRT Version: 8.6.3.1
GPU Type: GeForce RTX 4090
Nvidia Driver Version: 550.90.07
CUDA Version: 12.4
CUDNN Version: 12.3
Operating System + Version: Ubuntu 22.04.3 LTS
Python Version (if applicable): 3.10.12

Hi @ayagoz14finemax ,
can you please help us with the model and repro details?
thanks