Spectral Normalisation for ConvND Layers

Hi, I'm trying to do Spectral Norm for 2D and 3D convolutional layers. I found the code for a spectral normaliser in this notebook but it isn't quite working for me. The code I'm trying to run is as follows:

class SNConv2D(snt.Conv2D):
  """2D convolution with spectral normalisation."""

  def __init__(self,sn_eps=0.0001, *args,**kwargs):
    """Constructor."""
    super().__init__(*args,**kwargs)
    self._sn_eps = sn_eps
    self.spectral_normalizer = SpectralNormalizer()

  def __call__(self, inputs, is_training=True):
    self._initialize(inputs)
    self.spectral_normalizer._initialize(self.w)
    normed_w = self.spectral_normalizer(self.w, is_training=is_training)
    outputs = tf.matmul(inputs, normed_w)
    if self.with_bias:
      outputs = tf.add(outputs, self.b)
    return outputs

The spectral normaliser class is taken directly from the jupyter notebook. When I try to run this, the tf.matmul function throws the following error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: In[0] and In[1] must have compatible batch dimensions: >>[16,32,32,4] vs. [3,3,4,4] [Op:BatchMatMulV2]

[16,32,32,4] is the dimension of the input (which is correct) but [3,3,4,4], the output from the normaliser, doesn't seem right. Any help would be appreciated!

Edit: I think I see the issue. The code in the notebooks was written with the idea it would take 2D input (batch * length), like noise in a basic GAN, but I am trying to use it for a conditional GAN, where the input is of shape (batch * length * width * channels). How would one go about using spectral norm in that situation and aligning the shapes?