Spaces:
Sleeping
Sleeping
Update model.py
Browse files
model.py
CHANGED
@@ -7,12 +7,10 @@ from torch import nn
|
|
7 |
def create_effnetb2_model(num_classes:int=3,
|
8 |
seed:int=42):
|
9 |
"""Creates an EfficientNetB2 feature extractor model and transforms.
|
10 |
-
|
11 |
Args:
|
12 |
num_classes (int, optional): number of classes in the classifier head.
|
13 |
Defaults to 3.
|
14 |
seed (int, optional): random seed value. Defaults to 42.
|
15 |
-
|
16 |
Returns:
|
17 |
model (torch.nn.Module): EffNetB2 feature extractor model.
|
18 |
transforms (torchvision.transforms): EffNetB2 image transforms.
|
|
|
7 |
def create_effnetb2_model(num_classes:int=3,
|
8 |
seed:int=42):
|
9 |
"""Creates an EfficientNetB2 feature extractor model and transforms.
|
|
|
10 |
Args:
|
11 |
num_classes (int, optional): number of classes in the classifier head.
|
12 |
Defaults to 3.
|
13 |
seed (int, optional): random seed value. Defaults to 42.
|
|
|
14 |
Returns:
|
15 |
model (torch.nn.Module): EffNetB2 feature extractor model.
|
16 |
transforms (torchvision.transforms): EffNetB2 image transforms.
|