mirrors/mmengine
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 21:54:44 +08:00
Code Issues Projects Releases Wiki Activity

[Fix] Fix the incorrect device of inputs in get_model_complexity_info (#1130)

Browse Source
This commit is contained in:
CescMessi 2023-05-06 10:55:17 +08:00 committed by GitHub
parent 2085046d22
commit 83d76abc7f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mmengine/analysis/print_helper.py
View File

@ -725,14 +725,15 @@ def get_model_complexity_info(
raise ValueError('"input_shape" and "inputs" cannot be both set.') raise ValueError('"input_shape" and "inputs" cannot be both set.')
if inputs is None: if inputs is None:
device = next(model.parameters()).device
if is_tuple_of(input_shape, int): # tuple of int, construct one tensor if is_tuple_of(input_shape, int): # tuple of int, construct one tensor
inputs = (torch.randn(1, *input_shape), ) inputs = (torch.randn(1, *input_shape).to(device), )
elif is_tuple_of(input_shape, tuple) and all([ elif is_tuple_of(input_shape, tuple) and all([
is_tuple_of(one_input_shape, int) is_tuple_of(one_input_shape, int)
for one_input_shape in input_shape # type: ignore for one_input_shape in input_shape # type: ignore
]): # tuple of tuple of int, construct multiple tensors ]): # tuple of tuple of int, construct multiple tensors
inputs = tuple([ inputs = tuple([
torch.randn(1, *one_input_shape) torch.randn(1, *one_input_shape).to(device)
for one_input_shape in input_shape # type: ignore for one_input_shape in input_shape # type: ignore
]) ])
else: else: