Hi, code is too much, I just try in my code, seems problem occurred in torch.matmul() function.
below is my training process
FP32 each batch
Time: 0.00418s
Time: 0.00555s
Time: 0.00513s
Time: 0.00510s
Time: 0.00475s
Time: 0.00477s
Time: 0.00465s
Time: 0.00469s
in case of FP16 each batch:
Time: 0.46559s
Time: 0.46569s
Time: 0.46576s
Time: 0.46529s
Time: 0.46545s
Time: 0.46568s
Time: 0.46566s
Time: 0.46571s
Time: 0.46584s
my GPU is GTX1080ti