upgrade base image to pytorch 21.02

[wyli]

Is your feature request related to a problem? Please describe.
the dockerfile should also be flexible by accepting additional arguments when building an image. cc @madil90