diff --git a/Comparing TF and PT models.ipynb b/Comparing TF and PT models.ipynb
index 18b18a80..a7be90f7 100644
--- a/Comparing TF and PT models.ipynb	
+++ b/Comparing TF and PT models.ipynb	
@@ -1015,7 +1015,7 @@
     "    print(input_mask)\n",
     "    print(example_indices)\n",
     "    input_ids = input_ids.to(device)\n",
-    "    input_mask = input_mask.float().to(device)\n",
+    "    input_mask = input_mask.to(device)\n",
     "\n",
     "    all_encoder_layers, _ = model(input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)\n",
     "\n",