ReacherTN commited on
Commit
60d40f7
·
1 Parent(s): f29f05e

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +1 -1
model.py CHANGED
@@ -16,7 +16,7 @@ class Backbone(nn.Module) :
16
  class CustomModel(nn.Module) :
17
  def __init__(self) :
18
  super(CustomModel,self).__init__()
19
- self.backbone = Backbone("tf_efficientnetv2_b0",False)
20
  self.pooling = nn.AdaptiveAvgPool2d(1)
21
  self.head = nn.Linear(self.backbone.out_features,1)
22
  def forward(self,x) :
 
16
  class CustomModel(nn.Module) :
17
  def __init__(self) :
18
  super(CustomModel,self).__init__()
19
+ self.backbone = Backbone("tf_efficientnetv2_s",False)
20
  self.pooling = nn.AdaptiveAvgPool2d(1)
21
  self.head = nn.Linear(self.backbone.out_features,1)
22
  def forward(self,x) :