Skip to content

GitLab

Unverified Commit b2225436 authored by Bill Wu's avatar Bill Wu Committed by GitHub
Browse files

flops counter formatting fix (#3837)

parent 91ef554d
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
nni/compression/pytorch/utils/counter.py
View file @ b2225436
...@@ -313,12 +313,13 @@ class ModelProfiler: ...@@ -313,12 +313,13 @@ class ModelProfiler:
table.field_names = headers table.field_names = headers
for i, result in enumerate(self.results): for i, result in enumerate(self.results):
flops_count = int(result['flops'].item()) if isinstance(result['flops'], torch.Tensor) else int(result['flops'])
row_values = [ row_values = [
i, i,
result['name'], result['name'],
result['module_type'], result['module_type'],
str(result['weight_shape']), str(result['weight_shape']),
result['flops'], flops_count,
result['params'], result['params'],
] ]
name_counter[result['name']] += 1 name_counter[result['name']] += 1
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment