diff --git a/index.mjs b/index.mjs
index 56e44c8..69a335b 100755
--- a/index.mjs
+++ b/index.mjs
@@ -90,6 +90,7 @@ const model = await getModel(args.model)
 // If user has provided a model version, use it. Otherwise, use the latest version
 const modelVersionRegexp = /.*:[a-fA-F0-9]{64}$/
 const modelNameWithVersion = args.model.match(modelVersionRegexp) ? args.model : getModelNameWithVersion(model)
+const version = modelNameWithVersion.split(':')[1]
 
 const inputs = getModelInputs(model)
 
@@ -98,6 +99,7 @@ const indexFile = path.join(targetDir, 'index.js')
 const indexFileContents = fs.readFileSync(indexFile, 'utf8')
 const newContents = indexFileContents
   .replace('{{MODEL}}', modelNameWithVersion)
+  .replace('{{VERSION}}', version)
   .replace('\'{{INPUTS}}\'', JSON5.stringify(inputs, null, 2))
 fs.writeFileSync(indexFile, newContents)
 
diff --git a/template/index.js b/template/index.js
index 44168cc..99dce77 100644
--- a/template/index.js
+++ b/template/index.js
@@ -7,9 +7,14 @@ const replicate = new Replicate({
   userAgent: 'https://www.npmjs.com/package/create-replicate'
 })
 const model = '{{MODEL}}'
+const version = '{{VERSION}}'
 const input = '{{INPUTS}}'
 
 console.log({ model, input })
 console.log('Running...')
-const output = await replicate.run(model, { input })
-console.log('Done!', output)
+let prediction = await replicate.predictions.create({ version, input })
+prediction = await replicate.wait(prediction)
+
+console.log('Done!')
+console.log('Output:', prediction.output)
+console.log('View this prediction on the web:', `https://replicate.com/p/${prediction.id}`)