mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-12 07:10:52 +08:00
10 lines
578 B
Python
10 lines
578 B
Python
def append_dims(x, target_dims):
|
|
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
|
|
dims_to_append = target_dims - x.ndim
|
|
if dims_to_append < 0:
|
|
raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
|
|
expanded = x[(...,) + (None,) * dims_to_append]
|
|
# MPS will get inf values if it tries to index into the new axes, but detaching fixes this.
|
|
# https://github.com/pytorch/pytorch/issues/84364
|
|
return expanded.detach().clone() if expanded.device.type == 'mps' else expanded
|