fix bug reported by Shubham Tulsiani
add one-off function for testing with the built-in softmax