Dimension mismatch

[tianshunhan]

Thanks for your work！

I ran your code, but it has some problems. The reason is that the dimensions of ground-truth and prediction is not match. The error is as follows:

InvalidArgumentError:  Incompatible shapes: [32,257,6] vs. [32,9,18]
	 [[node my_loss_fn/sub (defined at tmp/ipykernel_7131/1937130251.py:4) ]] [Op:__inference_train_function_66805]

Errors may have originated from an input operation.
Input Source operations connected to node my_loss_fn/sub:
 IteratorGetNext (defined at tmp/ipykernel_7131/2605338104.py:5)

Function call stack:
train_function

By the way, how do you construct ground-truth for input data?
Could you update the latest code?