Update README.md

Co-authored-by: david-waterworth

README.md

@@ -32,5 +32,5 @@ predictions = torch.round((torch.sign(discriminator_outputs[0]) + 1) / 2)
 
 [print("%7s" % token, end="") for token in fake_tokens]
 
-[print("%7s" % int(prediction), end="") for prediction in predictions.tolist()]
+[print("%7s" % int(prediction), end="") for prediction in predictions.squeeze().tolist()]
 ```