PrathKothari commited on
Commit
d828e0f
·
verified ·
1 Parent(s): d2d4f6d

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +2 -0
model.py CHANGED
@@ -1,3 +1,5 @@
 
 
1
  def create_effnetb2_feature_extractor(num_classes: int, seed:int=42):
2
  weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
3
  transforms = weights.transforms()
 
1
+ import torch
2
+ import torchvision
3
  def create_effnetb2_feature_extractor(num_classes: int, seed:int=42):
4
  weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
5
  transforms = weights.transforms()