MolecularDiffusion.utils.torch¶

Attributes¶

`MAX_SEED_VALUE`
`MIN_SEED_VALUE`

Functions¶

`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.
`move_stray_tensor_attrs`(→ None)	Move plain (non-parameter, non-buffer) tensor attributes to device.
`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.

Module Contents¶

MolecularDiffusion.utils.torch.adjust_bias(param, new_shape)¶

MolecularDiffusion.utils.torch.adjust_weights(param, new_shape)¶

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.get_vram_size()¶

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

MolecularDiffusion.utils.torch.move_stray_tensor_attrs(module: torch.nn.Module, device: torch.device) → None¶

Move plain (non-parameter, non-buffer) tensor attributes to device.

nn.Module.to() only recurses into registered parameters/buffers/ submodules – a submodule that lazily caches a tensor as a plain self.foo = tensor attribute (rather than via register_buffer) is invisible to it, and stays wherever it was when the model was loaded. A checkpoint loaded with map_location=”cpu” remaps every tensor in the pickle regardless of whether it is registered, so such caches are stranded on CPU even after the model itself is moved to GPU. Call this right after moving a loaded model to its target device.

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.

MolecularDiffusion.utils.torch.MAX_SEED_VALUE = 4294967295¶

MolecularDiffusion.utils.torch.MIN_SEED_VALUE = 0¶