fix

main.py

@@ -383,7 +383,7 @@ def train(trainloader, optimizer, model, epoch):
         # print(loss.shape)
         # print(logits.shape)
         # exit()
-        loss = loss.sum()
+        loss = loss.mean()
         loss.backward()
         
         optimizer.step()

run.sh

-python main.py --dataset "RAVEN" --dataset_path "./dataset/I-RAVEN" -t center_single left_right up_down in_out distribute_four distribute_nine in_distri --use-gpu -v -lr 0.005 -vlr 0.005 -wd 0.01 -tsd 80 -fg 10 -nf 80 -chd 16 16 32 32 -rb -erb -vhd 128 -hd 64 32 -lhd 128 -bs 256
\ No newline at end of file
+python main.py --dataset "RAVEN" --dataset_path "./dataset/I-RAVEN" -t center_single left_right up_down in_out distribute_four distribute_nine in_distri --use-gpu -v -lr 0.005 -vlr 0.005 -wd 0.01 -tsd 80 -fg 10 -nf 80 -chd 16 16 32 32 -rb -erb -vhd 128 -hd 64 32 -lhd 128 -bs 2048
\ No newline at end of file