Unverified Commit c76e8840 authored by YiYi Xu's avatar YiYi Xu Committed by GitHub
Browse files

update get_parameter_dtype (#9526)



* up

* Update src/diffusers/models/modeling_utils.py
Co-authored-by: default avatarAryan <aryan@huggingface.co>

---------
Co-authored-by: default avatarAryan <aryan@huggingface.co>
parent d9c96917
...@@ -93,24 +93,20 @@ def get_parameter_device(parameter: torch.nn.Module) -> torch.device: ...@@ -93,24 +93,20 @@ def get_parameter_device(parameter: torch.nn.Module) -> torch.device:
def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype: def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype:
try: try:
params = tuple(parameter.parameters()) return next(parameter.parameters()).dtype
if len(params) > 0:
return params[0].dtype
buffers = tuple(parameter.buffers())
if len(buffers) > 0:
return buffers[0].dtype
except StopIteration: except StopIteration:
# For torch.nn.DataParallel compatibility in PyTorch 1.5 try:
return next(parameter.buffers()).dtype
def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]: except StopIteration:
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] # For torch.nn.DataParallel compatibility in PyTorch 1.5
return tuples
def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
gen = parameter._named_members(get_members_fn=find_tensor_attributes) tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
first_tuple = next(gen) return tuples
return first_tuple[1].dtype
gen = parameter._named_members(get_members_fn=find_tensor_attributes)
first_tuple = next(gen)
return first_tuple[1].dtype
class ModelMixin(torch.nn.Module, PushToHubMixin): class ModelMixin(torch.nn.Module, PushToHubMixin):
......
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