Torch max function update

[gulvarol]

Hi Alejandro,

I guess you updated the training code to work with the updated 'max' function which doesn't return index anymore. (with this commit)

But it doesn't seem to be updated in this demo repository. I haven't ran it myself, I just realized looking at the code.

[anewell]

Thanks for pointing that out, I think I didn't change it here since the bug only affected the on-the-fly accuracy calculation which doesn't get computed here. I've updated it to keep things consistent.

[anibali]

@anewell I think that your update addressing this issue is affecting the results adversely. See #14.