LukeOLuck commited on
Commit
b1449b3
1 Parent(s): f170aa7

fix indent

Browse files
Files changed (1) hide show
  1. model.py +8 -8
model.py CHANGED
@@ -3,15 +3,15 @@ import torchvision
3
  from torch import nn
4
 
5
  def create_efficientnet(output_shape: int):
6
- weights = torchvision.models.EfficientNet_V2_L_Weights.IMAGENET1K_V1.DEFAULT
7
- model = torchvision.models.efficientnet_v2_l(weights=weights)
8
 
9
- for param in model.parameters():
10
- param.requires_grad = False
11
 
12
- model.classifier = nn.Sequential(
13
- nn.Dropout(p=0.10, inplace=True),
14
- nn.Linear(in_features=1280, out_features=output_shape)
15
- )
16
 
17
  return model, weights.transforms()
 
3
  from torch import nn
4
 
5
  def create_efficientnet(output_shape: int):
6
+ weights = torchvision.models.EfficientNet_V2_L_Weights.IMAGENET1K_V1.DEFAULT
7
+ model = torchvision.models.efficientnet_v2_l(weights=weights)
8
 
9
+ for param in model.parameters():
10
+ param.requires_grad = False
11
 
12
+ model.classifier = nn.Sequential(
13
+ nn.Dropout(p=0.10, inplace=True),
14
+ nn.Linear(in_features=1280, out_features=output_shape)
15
+ )
16
 
17
  return model, weights.transforms()