Fix(Components/pytorch): Add single docker file for cpu and gpu

samples/contrib/pytorch-samples/Dockerfile

@@ -0,0 +1,26 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+ARG BASE_IMAGE=pytorch/pytorch:latest
+
+FROM ${BASE_IMAGE}
+
+COPY . .
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+RUN pip install --no-cache-dir pytorch-kfp-components
+
+ENV PYTHONPATH /workspace
+
+ENTRYPOINT /bin/bash

samples/contrib/pytorch-samples/README.md

@@ -37,13 +37,13 @@ Ex: `pipelines/samples/contrib/pytorch-samples/iris`
 
 ### Build and push the docker image
 ```
-docker build -t image_name:tag -f Dockerfile-cpu .
+docker build -t image_name:tag .
 ```
 
 to run the example in gpu, run the following commands for building docker image
 
 ```
-docker build -t image_name:tag -f Dockerfile-gpu .
+docker build --build-arg BASE_IMAGE=pytorch/pytorch:1.8.1-cuda10.2-cudnn7-runtime -t image_name:tag .
 ```
 
 push the docker image

samples/contrib/pytorch-samples/common/tensorboard/Dockerfile

@@ -0,0 +1,5 @@
+FROM gcr.io/deeplearning-platform-release/tf2-cpu.2-2:latest
+
+RUN pip install --user --no-cache-dir boto3
+
+RUN pip install --user --no-cache-dir torch-tb-profiler