How to convert this keras code to pytorch lightning?

I am trying to convert this keras code to pytorch lightning and use it. But some features are not there in pytorch to directly convert it. can someone help me to convert this to pytorch.

def sample(input_shape):
  input_layer = Input(shape=input_shape,name="INPUT")

  x = Conv3D(32, (3, 3, 3), strides=(1, 1, 1), padding="same", kernel_initializer=kernal_initializer)(input_layer)
  x = BatchNormalization()(x)
  x = Activation(activation)(x)
  x = MaxPooling3D(strides=2)(x)

  x = Conv3D(64, (3, 3, 3), strides=(1, 1, 1), padding="same", kernel_initializer=kernal_initializer)(x)
  x = BatchNormalization()(x)
  x = Activation(activation)(x)
  x = MaxPooling3D(strides=2)(x)

  x = Conv3D(128, (3, 3, 3), strides=(1, 1, 1), padding="same", kernel_initializer=kernal_initializer)(x)
  x = BatchNormalization()(x)
  x = Activation(activation)(x)
  x = MaxPooling3D(strides=2)(x)

  volumeSize = K.int_shape(x)
  x = Flatten()(x)

  x = Dense(256, activation=activation, kernel_initializer=tf.keras.initializers.glorot_normal)(x)

  x = Dense(128, activation=activation, kernel_initializer=tf.keras.initializers.glorot_normal, name="CODE")(x)

  x = Dense(256, activation=activation, kernel_initializer=tf.keras.initializers.glorot_normal)(x)

  x = Dense(np.prod(volumeSize[1:]), kernel_initializer=tf.keras.initializers.glorot_normal)(x)

  x = Reshape((volumeSize[1], volumeSize[2], volumeSize[3], volumeSize[4]))(x)

  x = Conv3DTranspose(128, (3, 3, 3), strides=(2, 2, 2), padding="same", kernel_initializer=kernal_initializer)(x)
  x = BatchNormalization()(x)
  x = Activation(activation)(x)

  x = Conv3DTranspose(64, (3, 3, 3), strides=(2, 2, 2), padding="same", kernel_initializer=kernal_initializer)(x)
  x = BatchNormalization()(x)
  x = Activation(activation)(x)

  x = Conv3DTranspose(32, (3, 3, 3), strides=(2, 2, 2), padding="same", kernel_initializer=kernal_initializer)(x)
  x = BatchNormalization()(x)
  x = Activation(activation)(x)

  x = Conv3D(1, (3, 3, 3), strides=(1, 1, 1),  padding="same", kernel_initializer=kernal_initializer, name = "OUTPUT")(x)
  x = Activation(tf.keras.activations.sigmoid)(x)

  autoencoder = Model(input_layer, x)
  autoencoder.summary()

  print("Model created successfully")
  return autoencoder

For Volume size i have experimented and found torch.stack(torch.unbind(Tensor)).numpy().shape Works fine. I am finding hard to convert this to pytorch.

My input Tensor size is (32, 55, 46,46) Here 32 is batch size, 55 channels, (46,46) is the image size