Hi all,
I am trying to use batchnorm1d layer in the linear output layer of my model as follows:
class output_layers(torch.nn.Module):
  def __init__(self, param):
    super(output_layers, self).__init__()
    self.param = param
    self.num_layers = self.param["decoder_num_layers"] 
    self.activation = self.param["activation_function"]
    self.input_size = self.param["lstm_hidden_size"] * self.param["seq_len"]
    self.factor = math.exp(np.log(1 / self.input_size) / self.num_layers)
    self.activation_layer = torch.nn.Mish()
   in_size = self.input_size
    for l in range(self.num_layers):
      if l != (self.num_layers - 1):
        self.decoder.add_module(f"decoder_block{l}", torch.nn.Sequential(torch.nn.Linear(in_size, round(in_size * self.factor)),
                                                                        self.activation_layer,
                                                                        torch.nn.BatchNorm1d(round(in_size * self.factor))))
        in_size = round(in_size*self.factor)
      else:
        self.decoder.add_module("output", torch.nn.Linear(in_size, 1))

  def forward(self, x):
    return self.decoder(x)

I plugged in this module into another main module with the forward implementation as below:
   def forward(self, x_con, x_cat):
    x = self.embedding(x_con, x_cat)
    x = x.view(-1, self.seq_len, self.input_size)
    if self.training == True:
      x = self.noise_layer(x)
    x = self.autoencoder(x)
    x = x.contiguous().view(-1, self.seq_len, self.encoded_final_size)
    output, (h_n, c_n) = self.lstm(x)
    output = output.contiguous().view(self.batch_size,-1)
    output = self.output_layer(output).squeeze()
    return output, (h_n, c_n)

While running though this output layer, the following error message pops out:
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/optuna/_optimize.py", line 216, in _run_trial
    value_or_values = func(trial)
  File "<ipython-input-18-78249d7efa36>", line 154, in __call__
    output, (h_n, c_n) = model(X_con, X_cat)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "<ipython-input-15-ab96b0392a9e>", line 35, in forward
    output = self.output_layer(output).squeeze()
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "<ipython-input-14-061a1e634053>", line 29, in forward
    return self.decoder(x)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/batchnorm.py", line 178, in forward
    self.eps,
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py", line 2282, in batch_norm
    input, weight, bias, running_mean, running_var, training, momentum, eps, torch.backends.cudnn.enabled
RuntimeError: CUDA error: an illegal memory access was encountered

The code was running with CUDA_LAUNCH_BLOCKING = 1
Cuda version is 11.0. GPU is Tesla P100 on google colab. Tried both pytorch version 1.9.0 and pytorch nighty. Same error message pops out. It only disappear if I deleted the batchnorm 1d layer.
Also tried to set “torch.backends.cudnn.enabled = False” but the error retain.
Most grateful if there are any solution to it. Much obliged.