Fix pytorch example (#1668)

Signed-off-by: Kevin Su <pingsutw@apache.org>

.github/workflows/checks.yml

@@ -53,7 +53,7 @@ jobs:
           # TODO: Register and update the examples below. (onnx_plugin, feast_integration, etc)
           echo "PACKAGES=$(find examples -maxdepth 1 -type d -exec basename '{}' \; \
           | grep -v -e 'testing' -e 'examples' \
-          | grep -v -e 'airflow_plugin' -e 'forecasting_sales' -e 'onnx_plugin' -e 'feast_integration' -e 'modin_plugin' -e 'sagemaker_inference_agent' -e 'kfpytorch_plugin' \
+          | grep -v -e 'airflow_plugin' -e 'forecasting_sales' -e 'onnx_plugin' -e 'feast_integration' -e 'modin_plugin' -e 'sagemaker_inference_agent' \
           | sort \
           | jq --raw-input . \
           | jq --slurp . \
@@ -72,8 +72,7 @@ jobs:
         example: "${{ fromJson(needs.list_examples.outputs.packages) }}"
     steps:
       - uses: actions/checkout@v4
-      - name: 'Clear action cache'
-        uses: ./.github/actions/clear-action-cache
+      - uses: insightsengineering/disk-space-reclaimer@v1
       - name: Set up Python
         uses: actions/setup-python@v4
         with:

examples/kfpytorch_plugin/kfpytorch_plugin/pytorch_mnist.py

@@ -34,17 +34,15 @@
 # :::
 #
 # The following imports are required to configure the PyTorch cluster in Flyte.
-# You can load them on demand.
 # %%
-if custom_image.is_container():
-    import matplotlib.pyplot as plt
-    import torch
-    import torch.nn.functional as F
-    from flytekitplugins.kfpytorch import PyTorch, Worker
-    from tensorboardX import SummaryWriter
-    from torch import distributed as dist
-    from torch import nn, optim
-    from torchvision import datasets, transforms
+import matplotlib.pyplot as plt
+import torch
+import torch.nn.functional as F
+from flytekitplugins.kfpytorch import PyTorch, Worker
+from tensorboardX import SummaryWriter
+from torch import distributed as dist
+from torch import nn, optim
+from torchvision import datasets, transforms
 
 # %% [markdown]
 # You can activate GPU support by either using the base image that includes the necessary GPU dependencies