diff --git a/docs/source-pytorch/starter/introduction.rst b/docs/source-pytorch/starter/introduction.rst
index 6f205dcd69..d1ffb8dbe0 100644
--- a/docs/source-pytorch/starter/introduction.rst
+++ b/docs/source-pytorch/starter/introduction.rst
@@ -204,7 +204,7 @@ Once you've trained the model you can export to onnx, torchscript and put it int
     encoder.eval()
 
     # embed 4 fake images!
-    fake_image_batch = Tensor(4, 28 * 28)
+    fake_image_batch = torch.rand(4, 28 * 28, device=autoencoder.device)
     embeddings = encoder(fake_image_batch)
     print("⚡" * 20, "\nPredictions (4 image embeddings):\n", embeddings, "\n", "⚡" * 20)