Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion httomo/runner/task_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,12 @@
ReadableDataSetSink,
)
from httomo.utils import save_2d_snapshot
from httomo.runner.gpu_utils import get_available_gpu_memory, gpumem_cleanup
from httomo.runner.gpu_utils import (
get_available_gpu_memory,
gpumem_cleanup,
gpu_enabled,
xp,
)
from httomo.runner.monitoring_interface import MonitoringInterface
from httomo.runner.pipeline import Pipeline
from httomo.runner.section import (
Expand Down Expand Up @@ -333,6 +338,11 @@ def _prepare(self):
self._check_params_for_sweep()
self._load_datasets()

if gpu_enabled:
xp.get_default_pinned_memory_pool().free_all_blocks()
xp.cuda.set_pinned_memory_allocator(None)
log_once("Disabled CuPy pinned memory pool", logging.DEBUG)

def _load_datasets(self):
start_time = self._log_task_start(
"loader",
Expand Down
Loading