Classification code runs into error.

[toshi2k2]

Command to reproduce-
python3 scripts/train.py --cate car --config config/pose_estimation_3d_nemo_cls.yaml --save_dir exp_cls/pose_estimation_3d_nemo_car

Error-

Traceback (most recent call last):
  File "scripts/train.py", line 104, in <module>
    main()
  File "scripts/train.py", line 100, in main
    train(cfg)
  File "scripts/train.py", line 71, in train
    loss_dict = model.train(sample)
  File "/home/NeMo/nemo/models/nemo_cls.py", line 212, in train
    get, y_idx, noise_sim, label_onehot = self.memory_bank(features, index, kpvis, label)
  File "/home/anaconda3/envs/nemo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/NeMo/nemo/models/feature_banks_cls.py", line 216, in forward
    get = torch.matmul(label_weight_onehot.transpose(0, 1), (x[:, 0:self.single_feature_dim, :] * visible.type(x.dtype).view(*visible.shape, 1)).view(x.shape[0], -1))
RuntimeError: The size of tensor a (1029) must match the size of tensor b (1024) at non-singleton dimension 1