Fixed same bug in another example notebook

Author: KevinMusgrave

examples/notebooks/DistributedTripletMarginLossMNIST.ipynb

@@ -171,6 +171,8 @@
         "def test(train_set, test_set, model, accuracy_calculator, data_device):\n",
         "    train_embeddings, train_labels = get_all_embeddings(train_set, model, data_device)\n",
         "    test_embeddings, test_labels = get_all_embeddings(test_set, model, data_device)\n",
+        "    train_labels = train_labels.squeeze(1)\n",
+        "    test_labels = test_labels.squeeze(1)\n",
         "    print(\"Computing accuracy\")\n",
         "    accuracies = accuracy_calculator.get_accuracy(test_embeddings, \n",
         "                                                train_embeddings,\n",