bevformer v2 can't be trained in the condition of batch size >= 2?

[TechnoNeko]

if i train in batch size of 2, error occurs as follow

Traceback (most recent call last):
  File "/home/yu/Project/BEVFormer/train_debug.py", line 265, in <module>
    main()
  File "/home/yu/Project/BEVFormer/train_debug.py", line 254, in main
    custom_train_model(
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/apis/train.py", line 27, in custom_train_model
    custom_train_detector(
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/apis/mmdet_train.py", line 199, in custom_train_detector
    runner.run(data_loaders, cfg.workflow)
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/mmcv/runner/epoch_based_runner.py", line 130, in run
    epoch_runner(data_loaders[i], **kwargs)
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/mmcv/runner/epoch_based_runner.py", line 51, in train
    self.run_iter(data_batch, train_mode=True, **kwargs)
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/mmcv/runner/epoch_based_runner.py", line 29, in run_iter
    outputs = self.model.train_step(data_batch, self.optimizer,
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/mmcv/parallel/data_parallel.py", line 75, in train_step
    return self.module.train_step(*inputs[0], **kwargs[0])
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/mmdet/models/detectors/base.py", line 248, in train_step
    losses = self(**data)
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/detectors/bevformerV2.py", line 161, in forward
    return self.forward_train(**kwargs)
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/detectors/bevformerV2.py", line 205, in forward_train
    prev_bev = self.obtain_history_bev(img_dict, prev_img_metas)
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/detectors/bevformerV2.py", line 179, in obtain_history_bev
    bev = self.pts_bbox_head(
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/dense_heads/bevformer_head.py", line 537, in forward
    return self.transformer.get_bev_features(
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/modules/transformerV2.py", line 240, in get_bev_features
    return super().forward(
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/modules/transformerV2.py", line 131, in forward
    bev_embed = self.encoder(bev_queries,
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/mmcv/runner/fp16_utils.py", line 110, in new_func
    return old_func(*args, **kwargs)
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/modules/encoder.py", line 193, in forward
    reference_points_cam, bev_mask = self.point_sampling(
  File "/home/yu/software/anaconda3/envs/lanesegnet/lib/python3.8/site-packages/mmcv/runner/fp16_utils.py", line 198, in new_func
    return old_func(*args, **kwargs)
  File "/home/yu/Project/BEVFormer/projects/mmdet3d_plugin/bevformer/modules/encoder.py", line 119, in point_sampling
    lidar2img = lidar2img.view(
RuntimeError: shape '[1, 2, 6, 1, 4, 4]' is invalid for input of size 96

Process finished with exit code 1