fix pred_visualization

src/basicts/models/Koopa/callback/koopa_mask_init.py

@@ -9,7 +9,7 @@
 
 
 class KoopaMaskInitCallback(BasicTSCallback):
-    
+
     """Callback for initializing Koopa mask during training.
 
         Changes made:

src/basicts/runners/basicts_runner.py

@@ -7,7 +7,6 @@
 from contextlib import nullcontext
 from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, Union
 
-import numpy as np
 import setproctitle
 import torch
 from easytorch.core.checkpoint import (backup_last_ckpt, clear_ckpt, load_ckpt,
@@ -16,6 +15,7 @@
 from easytorch.utils import (TimePredictor, get_local_rank, get_logger,
                              is_master, master_only)
 from easytorch.utils.env import get_rank, set_tf32_mode, setup_determinacy
+from numpy.lib.format import open_memmap
 from torch import nn
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data.distributed import DistributedSampler
@@ -35,6 +35,7 @@
     from basicts.configs import BasicTSConfig
 
 
+
 class BasicTSRunner:
     """
     A base runner that uses epoch as the fundamental training unit.
@@ -1082,11 +1083,11 @@ def _save_results(self, batch_idx: int, batch_data: Dict[str, torch.Tensor]) ->
 
         # create memmap files
         if batch_idx == 0:
-            self._inputs_memmap = np.memmap(inputs_path, dtype=inputs.dtype, mode="w+",
+            self._inputs_memmap = open_memmap(inputs_path, dtype=inputs.dtype, mode="w+",
                                     shape=(total_samples, *inputs.shape[1:]))
-            self._prediction_memmap = np.memmap(pred_path, dtype=prediction.dtype, mode="w+",
+            self._prediction_memmap = open_memmap(pred_path, dtype=prediction.dtype, mode="w+",
                                     shape=(total_samples, *prediction.shape[1:]))
-            self._targets_memmap = np.memmap(targets_path, dtype=targets.dtype, mode="w+",
+            self._targets_memmap = open_memmap(targets_path, dtype=targets.dtype, mode="w+",
                                 shape=(total_samples, *targets.shape[1:]))
 
         start = batch_idx * inputs.shape[0]