From df8982efd1911f327f9d0732cf41a97f7aa925a5 Mon Sep 17 00:00:00 2001 From: ipesanz Date: Mon, 14 Jun 2021 02:32:12 +0200 Subject: [PATCH] Update 15_transfer_learning.py Add comment to have a FC layer of the same size as len(class_name)... it was commented on l:138 but in l:172 --- 15_transfer_learning.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/15_transfer_learning.py b/15_transfer_learning.py index e68c4c85..60756f3b 100644 --- a/15_transfer_learning.py +++ b/15_transfer_learning.py @@ -167,7 +167,9 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25): # Parameters of newly constructed modules have requires_grad=True by default num_ftrs = model_conv.fc.in_features -model_conv.fc = nn.Linear(num_ftrs, 2) + +# Creating a FC layer of the same size as our categories +model_conv.fc = nn.Linear(num_ftrs, len(class_names)) model_conv = model_conv.to(device)