Fix import

TensorDot op does not exist anymore