MolecularDiffusion.utils.torch¶

Attributes¶

`MAX_SEED_VALUE`
`MIN_SEED_VALUE`

`adjust_bias`(param, new_shape)
`adjust_weights`(param, new_shape)
`cat`(objs, args, *kwargs)	Concatenate a list of nested containers with the same structure.
`clone`(obj, args, *kwargs)	Clone tensors in any nested conatiner.
`cpu`(obj, args, *kwargs)	Transfer any nested container of tensors to CPU.
`cuda`(obj, args, *kwargs)	Transfer any nested container of tensors to CUDA.
`detach`(obj)	Detach tensors in any nested conatiner.
`get_vram_size`()
`mean`(obj, args, *kwargs)	Compute mean of tensors in any nested container.
`recursive_module_to_device`(module, device)	Recursively assigns a given device to all submodules of a torch.nn.Module.
`seed_everything`(→ int)	Sets seed for reproducibility across torch, numpy, and random modules.
`stack`(objs, args, *kwargs)	Stack a list of nested containers with the same structure.

MolecularDiffusion.utils.torch.cat(objs, *args, **kwargs)¶: Concatenate a list of nested containers with the same structure.

MolecularDiffusion.utils.torch.clone(obj, *args, **kwargs)¶: Clone tensors in any nested conatiner.

MolecularDiffusion.utils.torch.cpu(obj, *args, **kwargs)¶: Transfer any nested container of tensors to CPU.

MolecularDiffusion.utils.torch.cuda(obj, *args, **kwargs)¶: Transfer any nested container of tensors to CUDA.

MolecularDiffusion.utils.torch.detach(obj)¶: Detach tensors in any nested conatiner.

MolecularDiffusion.utils.torch.mean(obj, *args, **kwargs)¶: Compute mean of tensors in any nested container.

MolecularDiffusion.utils.torch.recursive_module_to_device(module: torch.nn.Module, device: torch.device)¶

Recursively assigns a given device to all submodules of a torch.nn.Module.

Parameters:

module (nn.Module) – The main module to which the device needs to be assigned.
device (torch.device) – The target device (e.g., torch.device(‘cuda’) or torch.device(‘cpu’)).

MolecularDiffusion.utils.torch.seed_everything(seed: int = None, workers: bool = False, verbose: bool = True) → int¶

Sets seed for reproducibility across torch, numpy, and random modules.

Parameters:

seed (int, optional) – The seed to use. If None, it checks ‘PL_GLOBAL_SEED’ in env or defaults to 0.
workers (bool) – Whether to set the ‘PL_SEED_WORKERS’ env variable.
verbose (bool) – If True, logs the chosen seed.

MolecularDiffusion.utils.torch.stack(objs, *args, **kwargs)¶: Stack a list of nested containers with the same structure.