diff --git a/examples/data_types_and_io/data_types_and_io/pytorch_types.py b/examples/data_types_and_io/data_types_and_io/pytorch_types.py
index a3cd61180..3a034e73b 100644
--- a/examples/data_types_and_io/data_types_and_io/pytorch_types.py
+++ b/examples/data_types_and_io/data_types_and_io/pytorch_types.py
@@ -73,7 +73,6 @@ def pytorch_native_wf():
 # %% [markdown]
 # Passing around tensors and modules is no more a hassle!
 
-from dataclasses import dataclass
 
 # %% [markdown]
 # ## Checkpoint
@@ -85,6 +84,8 @@ def pytorch_native_wf():
 # As per PyTorch [docs](https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model), it is recommended to
 # store the module's `state_dict` rather than the module itself. However, the serialization should work either way.
 # %%
+from dataclasses import dataclass
+
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim