Thanks for the quick response! I tried that modification adding the threshold and permuting the output tensor indices and I'm still getting the same error. (I also tried all other permutations of the indices 0:1:2:3,50 / 3:2:1:0,50 / 3:1:2:0 / etc) Did you make any other changes to get it to work well?