Having Issue with GRU while trying to use DreamverV3

Hi,
I am having an issue with GRU layer in World_Model while trying to use DreamverV3, below is the error

ValueError: Exception encountered when calling GRUCell.call().
        
        Attempting to slice scalar input. for '{{node gru_1/gru_cell_1/strided_slice}} = StridedSlice[Index=DT_INT32, T=DT_INT32, begin_mask=0, ellipsis_mask=0, end_mask=0, new_axis_mask=0, shrink_axis_mask=0](gru_1/gru_cell_1/shape_as_tensor, gru_1/gru_cell_1/strided_slice/stack, gru_1/gru_cell_1/strided_slice/stack_1, gru_1/gru_cell_1/strided_slice/stack_2)' with input shapes: [], [1], [1], [1] and with computed input tensors: input[3] = <1>.
        
        Arguments received by GRUCell.call():
          • inputs=tf.Tensor(shape=(?, 1024), dtype=float32)
          • states=('tf.Tensor(shape=(?, 4096), dtype=float32)',)
          • training=False
    
    
    Arguments received by SequenceModel.call():
      • a=tf.Tensor(shape=(?, 3), dtype=float32)
      • h=tf.Tensor(shape=(?, 4096), dtype=float32)
      • z=tf.Tensor(shape=(?, 32, 32), dtype=float32)


Arguments received by DreamerModel.call():
  • inputs=None
  • observations=tf.Tensor(shape=(1, 64, 2), dtype=float32)
  • actions=tf.Tensor(shape=(1, 64, 3), dtype=float32)
  • is_first=tf.Tensor(shape=(1, 64), dtype=bool)
  • start_is_terminated_BxT=tf.Tensor(shape=(64,), dtype=bool)
  • gamma=0.997

Can someone help me how to solve this issue?