MemPoolContext

class torch.cuda.MemPoolContext(*args, **kwargs)

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()

Returns the active MemPool

Return type

Optional[_MemPool]