Translator.unload_model(to_cpu=True) takes long time about ctranslate2 HOT 8 CLOSED

translator = Translator(model_path, "cuda", compute_type) #  (takes ~6 seconds)
translator.unload_model(to_cpu=True) # (takes a few ms)
translator.load_model() # (takes ~14 seconds)

translator.unload_model(to_cpu=True) # (takes a few ms)
translator.load_model() # (takes ~1.5 seconds)

      void unload_model(const bool to_cpu) {
        if (to_cpu && _device == Device::CPU)
          return;

        // Do not unload the model if some batches are still being processed.
        if (_pool->num_active_batches() > 0)
          return;

        // If the lock is not acquired immediately it means the model is being used
        // in another thread and we can't unload it at this time.
        std::unique_lock lock(_mutex, std::try_to_lock);
        if (!lock || !_model_is_loaded)
          return;

        _gpu_cached_models = _pool->detach_models();
        if (!to_cpu)
          _cpu_cached_models.clear();
        else
          _gpu_cached_models.clear();

        // We clear the CUDA allocator cache to further reduce the memory after unloading the model.
        if (_device == Device::CUDA)
          _pool->clear_cache();

        _model_is_loaded = false;
      }

      void load_model() {
        std::unique_lock lock(_mutex);
        if (_model_is_loaded)
          return;

        if (_cpu_cached_models.empty()) {
          _cpu_cached_models = _model_loader.load();
          move_cached_models(Device::CPU, std::vector<int>(_cpu_cached_models.size(), 0), _cpu_cached_models);
          _gpu_cached_models = _model_loader.load();
        } else {
          _gpu_cached_models = copy_cached_models(_device, _device_index);
          move_cached_models(_device, _device_index, _gpu_cached_models, _num_replicas_per_device);
        }

        _pool->set_models(_gpu_cached_models);
        _gpu_cached_models.clear();
        _model_is_loaded = true;
      }

    private:
      const Device _device;
      const std::vector<int>& _device_index;
      const size_t _num_replicas_per_device;
      std::vector<std::shared_ptr<const models::Model>> _cpu_cached_models;
      std::vector<std::shared_ptr<const models::Model>> _gpu_cached_models;
      bool _model_is_loaded;

      // Use a shared mutex to protect the model state (loaded/unloaded).
      // Multiple threads can read the model at the same time, but a single thread can change
      // the model state (e.g. load or unload the model).
      std::shared_mutex _mutex;

      void assert_model_is_ready() const {
        if (!_model_is_loaded)
          throw std::runtime_error("The model for this translator was unloaded");
      }

      void move_cached_models(Device device,
                              const std::vector<int>& device_index,
                              std::vector<std::shared_ptr<const models::Model>> _cached_models,
                              size_t num_models_per_device = 1) {
        for (size_t i = 0; i < _cached_models.size(); ++i) {
          auto& model = const_cast<models::Model&>(*_cached_models[i]);
          model.set_device(device, device_index[i / num_models_per_device]);
        }
      }

      std::vector<std::shared_ptr<const models::Model>> copy_cached_models(Device device,
                                                                           const std::vector<int>& device_index,
                                                                           size_t num_models_per_device = 1) {
        std::vector<std::shared_ptr<const models::Model>> copied_models;
        for (size_t i = 0; i < _cpu_cached_models.size(); ++i) {
          auto& model = const_cast<models::Model&>(*_cpu_cached_models[i]);
          auto copied_model = model.copy_to(device, device_index[i / num_models_per_device]);
          copied_models.push_back(copied_model);
        }
        return copied_models;
      }
    };

      void load_model(bool keep_cache) {
        std::unique_lock lock(_mutex);
        if (_model_is_loaded)
          return;

        if (_cached_models.empty()) {
          _cached_models = _model_loader.load();
        } else {
          move_cached_models(_device, _device_index, _num_replicas_per_device);
        }

        _pool->set_models(_cached_models);
        if (!keep_cache)
          _cached_models.clear();
        _model_is_loaded = true;
      }

      void unload_model(const bool to_cpu) {
        if (to_cpu && _device == Device::CPU)
          return;

        // Do not unload the model if some batches are still being processed.
        if (_pool->num_active_batches() > 0)
          return;

        // If the lock is not acquired immediately it means the model is being used
        // in another thread and we can't unload it at this time.
        std::unique_lock lock(_mutex, std::try_to_lock);
        if (!lock || !_model_is_loaded)
          return;
       
        if (_cached_models.empty) {
          _cached_models = _pool->detach_models();
          if (to_cpu)
            move_cached_models(Device::CPU, std::vector<int>(_cached_models.size(), 0));
          else
            _cached_models.clear();
        }

        // We clear the CUDA allocator cache to further reduce the memory after unloading the model.
        if (_device == Device::CUDA)
          _pool->clear_cache();

        _model_is_loaded = false;
      }