Munzali commited on
Commit
a9c42e5
·
verified ·
1 Parent(s): 00bd9f5

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +2 -2
model.py CHANGED
@@ -20,7 +20,7 @@ def create_mobilenet_model(num_classes:int=4,
20
  """
21
  # Create EffNetB2 pretrained weights, transforms and model
22
 
23
- transforms = transforms.Compose([
24
  transforms.Resize((224, 224)), # 1. Reshape all images to 224x224 (though some models may require different sizes)
25
  transforms.ToTensor(), # 2. Turn image values to between 0 & 1
26
  transforms.Normalize(mean=[0.485, 0.456, 0.406], # 3. A mean of [0.485, 0.456, 0.406] (across each colour channel)
@@ -45,5 +45,5 @@ def create_mobilenet_model(num_classes:int=4,
45
  bias=True)
46
  )
47
 
48
- return model, transforms
49
 
 
20
  """
21
  # Create EffNetB2 pretrained weights, transforms and model
22
 
23
+ transform = transforms.Compose([
24
  transforms.Resize((224, 224)), # 1. Reshape all images to 224x224 (though some models may require different sizes)
25
  transforms.ToTensor(), # 2. Turn image values to between 0 & 1
26
  transforms.Normalize(mean=[0.485, 0.456, 0.406], # 3. A mean of [0.485, 0.456, 0.406] (across each colour channel)
 
45
  bias=True)
46
  )
47
 
48
+ return model, transform
49