mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Add drop args to benchmark.py

Browse Source
This commit is contained in:
Ross Wightman 2021-12-14 13:51:00 -08:00
parent ab49d275de
commit 683fba7686
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
benchmark.py
View File

@ -199,7 +199,11 @@ class BenchmarkRunner:
num_classes=kwargs.pop('num_classes', None), num_classes=kwargs.pop('num_classes', None),
in_chans=3, in_chans=3,
global_pool=kwargs.pop('gp', 'fast'), global_pool=kwargs.pop('gp', 'fast'),
scriptable=torchscript) scriptable=torchscript,
drop_rate=kwargs.pop('drop', 0.),
drop_path_rate=kwargs.pop('drop_path', None),
drop_block_rate=kwargs.pop('drop_block', None),
)
self.model.to( self.model.to(
device=self.device, device=self.device,
dtype=self.model_dtype, dtype=self.model_dtype,