Update README.md

README.md

@@ -28,12 +28,14 @@ if __name__ == '__main__':
         print(f'usage: {argv[0]} <model> <file 1>')
         exit(1)
 
+    device = "cpu"
+
     processor = Wav2Vec2Processor.from_pretrained(argv[1])
-    model = Wav2Vec2ForCTC.from_pretrained(argv[1])
+    model = Wav2Vec2ForCTC.from_pretrained(argv[1]).to(device)
 
     f = argv[2]
     s,sample_rate = sf.read(f)
-    input_values = processor(s, return_tensors="pt").input_values
+    input_values = processor(s, return_tensors="pt").to(device).input_values
     logits = model(input_values).logits
     predicted_ids = torch.argmax(logits, dim=-1)