torch_compile¶

keras_mml.utils.misc.coverage.torch_compile(model=None, **kwargs)[source]¶

Custom decorator similar to torch.compile(). However the compilation of the function will not be performed if the DISABLE_TORCH_COMPILE environment variable is set.

Parameters:

model (Optional[Callable], default: None) – Module/function to optimize.
*args – Arguments for torch.compile().
**kwargs – Keyword arguments for torch.compile().

Returns:

Callable – Decorated model.