I’d say it depends on the length of sequence_of_tensors.
are all the tensors in sequence of tensors the same length? You can try calculate that batch-wise rather than list-wise.
I would also store A-B as a variable, since you are performing that op 3 times I think it will create 3 graph nodes.
So, in short, I would try to do something like:
A,B,C =stacked_sequence_of_tensors:
tmp = (A-B)**2
loss = (tmp * k1 * (A > C).float() +
             tmp * k2 * (A < C).float() +
             tmp * k3 * (A == C).float()).sum()