采用版本pytorch:1.8.1。

在用timm库跑自定义的模型时，报了以下错误：

Traceback (most recent call last): 

  File "train.py", line 908, in <module> 

Traceback (most recent call last): 

  File "train.py", line 908, in <module> 

    main() 

  File "train.py", line 690, in main 

Traceback (most recent call last): 

  File "train.py", line 908, in <module> 

Traceback (most recent call last): 

  File "train.py", line 908, in <module> 

    amp_autocast=amp_autocast, loss_scaler=loss_scaler, model_ema=model_ema, mixup_fn=mixup_fn) 

  File "train.py", line 778, in train_one_epoch 

    main() 

  File "train.py", line 690, in main 

    loss.backward(create_graph=second_order) 

  File "/home/ma-user/anaconda3/envs/torch1.8.1/lib/python3.7/site-packages/torch/tensor.py", line 245, in backward 

    amp_autocast=amp_autocast, loss_scaler=loss_scaler, model_ema=model_ema, mixup_fn=mixup_fn) 

  File "train.py", line 778, in train_one_epoch 

    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs) 

  File "/home/ma-user/anaconda3/envs/torch1.8.1/lib/python3.7/site-packages/torch/autograd/__init__.py", line 147, in backward 

    main() 

  File "train.py", line 690, in main 

    allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag 

RuntimeError: The size of tensor a (320) must match the size of tensor b (4) at non-singleton dimension 2 

求救怎么解决，急求，谢谢各位大佬。本来报的错是cuda not enabled，我是改了相应的distributed.py文件中的torch.cuda.xxx全改成torch.npu.xxx了，现在报错是张量不匹配，但是我这个模型在GPU上可以正常跑通。