Spaces:

ishaal007
/

gadgets_classifier

Sleeping

ishaal007 commited on Oct 13, 2022

Commit

5ba68b9

1 Parent(s): f6e125b

Update model.py

Files changed (1) hide show

model.py CHANGED Viewed

@@ -20,6 +20,6 @@ def create_gadgets_model(num_classes:int=3,
     model.fc = nn.Sequential(
                nn.Linear(2048, 128),
                nn.ReLU(inplace=True),
-               nn.Linear(in_features= 128,out_features=output_shape))
     return model, transforms

     model.fc = nn.Sequential(
                nn.Linear(2048, 128),
                nn.ReLU(inplace=True),
+               nn.Linear(in_features= 128,out_features=num_classes))
     return model, transforms