Small fix in engine test
Small fix to logits test