MemPoolContext — PyTorch 2.7 documentation (original) (raw)

class torch.cuda.MemPoolContext(*args, **kwargs)[source][source]¶

MemPoolContext holds the currently active pool and stashes the previous pool. On deletion it makes the previous pool active.

Parameters

• pool (torch.cuda.MemPool) – a MemPool object to be made active so that
• pool. (allocations route to this) –

static active_pool()[source][source]¶

Returns the active MemPool

Return type

Optional[__MemPool_]