bentrevett
diff --git a/‎4 - VGG.ipynb
Lines changed: 3 additions & 1 deletion b/‎4 - VGG.ipynb
Lines changed: 3 additions & 1 deletion
@@ -595,7 +595,9 @@
    },
    "outputs": [],
    "source": [
-    "final_fc = nn.Linear(4096, 10)"
+    "IN_FEATURES = pretrained_model.classifier[-1].in_features \n",
+    "\n",
+    "final_fc = nn.Linear(IN_FEATURES, OUTPUT_DIM)"
    ]
   },
   {
Original file line number	Diff line number	Diff line change
`@@ -595,7 +595,9 @@`
`595`	`595`	`},`
`596`	`596`	`"outputs": [],`
`597`	`597`	`"source": [`
`598`		`- "final_fc = nn.Linear(4096, 10)"`
	`598`	`+ "IN_FEATURES = pretrained_model.classifier[-1].in_features \n",`
	`599`	`+ "\n",`
	`600`	`+ "final_fc = nn.Linear(IN_FEATURES, OUTPUT_DIM)"`
`599`	`601`	`]`
`600`	`602`	`},`
`601`	`603`	`{`