Support for arm64 in pytorch-mnist image

[tenzen-y]

What this PR does / why we need it:
I modified Dockerfile for pytorch-mnist to support arm64.
Also, I have tested using the following commands:

$ # For linux/arm64 with cuda
$ docker run -it --pull always --platform linux/arm64 ytenzen/multiplatform-pytorch-gpu:latest --epochs=1 --batch-size=256
$ # For linux/amd64 with cuda
$ docker run -it --pull always --platform linux/amd64 ytenzen/multiplatform-pytorch-gpu:latest --epochs=1 --batch-size=256
$ # For linux/arm64
$ docker run -it --pull always --platform linux/arm64 ytenzen/multiplatform-pytorch-cpu:latest --epochs=1 --batch-size=256
$ # For linux/amd64
$ docker run -it --pull always --platform linux/amd64 ytenzen/multiplatform-pytorch-cpu:latest --epochs=1 --batch-size=256

/hold until #1940 is merged

Which issue(s) this PR fixes (optional, in fixes #<issue number>(, fixes #<issue_number>, ...) format, will close the issue(s) when PR gets merged):
Part-of #1900

Checklist:

• Docs included if any changes are user facing

[coveralls]

Coverage increased (+0.1%) to 73.545% when pulling fced2a3 on tenzen-y:support-for-arm64-in-pytorch-mnist-image into 3906f8e on kubeflow:master.

[tenzen-y]

/hold cancel

[tenzen-y]

This PR is ready to merge.

/assign @johnugeorge @gaocegege @andreyvelich

[johnugeorge]

/lgtm
/approve

[google-oss-prow]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: johnugeorge, tenzen-y

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

• OWNERS [johnugeorge]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment