mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
@@ -50,7 +50,7 @@ def hidream_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
|
||||
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
|
||||
omega = 1.0 / (theta**scale)
|
||||
|
||||
batch_size, _seq_length = pos.shape
|
||||
batch_size, seq_length = pos.shape
|
||||
out = torch.einsum("...n,d->...nd", pos, omega)
|
||||
cos_out = torch.cos(out)
|
||||
sin_out = torch.sin(out)
|
||||
|
||||
@@ -52,7 +52,7 @@ def autocast_init(self, device_type=None, dtype=None, enabled=True, cache_enable
|
||||
|
||||
original_grad_scaler_init = torch.amp.grad_scaler.GradScaler.__init__
|
||||
@wraps(torch.amp.grad_scaler.GradScaler.__init__)
|
||||
def GradScaler_init(self, device: str | None = None, init_scale: float = 2.0**16, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True):
|
||||
def GradScaler_init(self, device: str = None, init_scale: float = 2.0**16, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True):
|
||||
if device is None or check_cuda(device):
|
||||
return original_grad_scaler_init(self, device=return_xpu(device), init_scale=init_scale, growth_factor=growth_factor, backoff_factor=backoff_factor, growth_interval=growth_interval, enabled=enabled)
|
||||
else:
|
||||
|
||||
Reference in New Issue
Block a user