* align squeeze module with torch

* fix comment

* fix argmax bug

* fix bug

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>