"vscode:/vscode.git/clone" did not exist on "51472e756a1e43cc4e6d36597920f382af673d71"
Unverified Commit d8420f81 authored by Boyuan Yao's avatar Boyuan Yao Committed by GitHub
Browse files

[hotfix] fix wrong type name in profiler (#1678)

parent 132b4306
from functools import partial
from typing import Callable, Any, Dict, Tuple
import torch
from torch.nn.parameter import Parameter
from torch.fx import Graph, Node
from torch.fx.node import Argument, Target
from torch.utils._pytree import tree_map
......@@ -298,7 +299,7 @@ def profile_function(target: 'Target', device: str = 'meta') -> Callable:
param_size = 0
def get_param_size(x):
if isinstance(x, torch.nn.parameter):
if isinstance(x, Parameter):
param_size += activation_size(x)
tree_map(get_param_size, args)
......
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