mirror of
https://github.com/vladmandic/automatic
synced 2026-08-29 00:20:59 +02:00
Vladimir Mandic
20 lines
513 B
Python
20 lines
513 B
Python
import torch
|
|
from collections.abc import Callable
|
|
from modules.shared import log, opts
|
|
|
|
|
|
def catch_nan(func: Callable[[], torch.Tensor]):
|
|
if not opts.directml_catch_nan:
|
|
return func()
|
|
|
|
tries = 0
|
|
tensor = func()
|
|
while tensor.isnan().sum() != 0 and tries < 10:
|
|
if tries == 0:
|
|
log.warning("NaN is produced. Retry with same values...")
|
|
tries += 1
|
|
tensor = func()
|
|
if tensor.isnan().sum() != 0:
|
|
log.error("Failed to cover NaN.")
|
|
return tensor
|