magrittetorch.utils.memorymapping

Contains a simple class to help with GPU memory management.

class magrittetorch.utils.memorymapping.MemoryManager

Simple memory management object. Handles memory issues by splitting the input data when neccessary. For now, it is assumed that the user only has one GPU to compute on. TODO: add support for multiple GPU’s.

run_with_splitting(lambdafun: Callable[[Tensor], Any], data: Tensor, device: device) → Tensor | List[Tensor]

Runs the given function with the given input, and if an out of memory error is thrown, halves the input and tries again

Parameters:

• lambdafun (callable[[torch.Tensor], Any]) – function to run

• data (torch.Tensor) – input data, which can be split in the first dimension if an out of memory error is thrown