some trouble about SVX loss of pytorch

i code the SVX loss of pytorch , like this:
-------------------------------------------------------------
def forward(self, cos_theta, label):
cos_theta = cos_theta.clamp(-1, 1) # for numerical stability
batch_size = label.size(0)
gt = cos_theta[torch.arange(0, batch_size), label].view(-1, 1) # ground truth score
index = torch.where(label != -1)[0]
print('index shape', cos_theta.shape)
if self.xtype == 'MV-AM':
mask = cos_theta > gt - self.m
hard_vector = cos_theta[mask]
cos_theta[mask] = (self.t + 1.0) * hard_vector + self.t # adaptive
final_gt = gt - self.m
elif self.xtype == 'MV-Arc':
sin_theta = torch.sqrt(1.0 - torch.pow(gt, 2))
cos_theta_m = gt * self.cos_m - sin_theta * self.sin_m # cos(gt + margin)
mask = cos_theta > cos_theta_m
hard_vector = cos_theta[mask]
cos_theta[mask] = (self.t + 1.0) * hard_vector + self.t # adaptive
final_gt = cos_theta_m
# final_gt = torch.where(gt > cos_theta_m, cos_theta_m, gt)
else:
raise Exception('unknown xtype!')

#cos_theta.scatter_(1, label.data.view(-1, 1), final_gt)
cos_theta.scatter_(1, label[index,None], final_gt)
cos_theta *= self.s
return cos_theta
------------------------------------------------------
Traceback (most recent call last):
File "partial_fc.py", line 291, in
main(args.local_rank)
File "partial_fc.py", line 235, in main
logits.backward(grad)
File "/home/lthpc/.conda/envs/partial_fc/lib/python3.6/site-packages/torch/tensor.py", line 185, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph)
File "/home/lthpc/.conda/envs/partial_fc/lib/python3.6/site-packages/torch/autograd/__init__.py", line 127, in backward
allow_unreachable=True) # allow_unreachable flag
RuntimeError: Function ScatterBackward0 returned an invalid gradient at index 1 - got [29, 1] but expected shape compatible with [256, 1]
Exception raised from validate_outputs at /pytorch/torch/csrc/autograd/engine.cpp:602 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x42 (0x7f7cdfd231e2 in /home/lthpc/.conda/envs/partial_fc/lib/python3.6/site-packages/torch/lib/libc10.so)
frame #1: + 0x336b180 (0x7f7d1a0e3180 in /home/lthpc/.conda/envs/partial_fc/lib/python3.6/site-packages/torch/lib/libtorch_cpu.so)
frame #2: torch::autograd::Engine::evaluate_function(std::shared_ptr&, torch::autograd::Node*, torch::autograd::InputBuffer&, std::shared_ptr const&) + 0x3fd (0x7f7d1a0e83fd in /home/lthpc/.conda/envs/partial_fc/lib/python3.6/site-packages/torch/lib/libtorch_cpu.so)
frame #3: torch::autograd::Engine::thread_main(std::shared_ptr const&) + 0x451 (0x7f7d1a0e9fa1 in /home/lthpc/.conda/envs/partial_fc/lib/python3.6/site-packages/torch/lib/libtorch_cpu.so)
frame #4: torch::autograd::Engine::thread_init(int, std::shared_ptr const&, bool) + 0x89 (0x7f7d1a0e2119 in /home/lthpc/.conda/envs/partial_fc/lib/python3.6/site-packages/torch/lib/libtorch_cpu.so)
frame #5: torch::autograd::python::PythonEngine::thread_init(int, std::shared_ptr const&, bool) + 0x4a (0x7f7d2788234a in /home/lthpc/.conda/envs/partial_fc/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
frame #6: + 0xb8408 (0x7f7d287d4408 in /opt/anaconda3/lib/libstdc++.so.6)
frame #7: + 0x76ba (0x7f7d29e346ba in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #8: clone + 0x6d (0x7f7d29b6a4dd in /lib/x86_64-linux-gnu/libc.so.6)
---------------------------------------------------------------------------------------

thanks