code: def predict(
      model: tf.keras.Model | torch.nn.Module, x_batch: np.ndarray
  ) -> np.ndarray:
      if isinstance(model, torch.nn.Module):
          with torch.no_grad():
              return model(torch.Tensor(x_batch)).argmax(axis=1).numpy()
      else:
          return model.predict(x_batch, verbose=0).argmax(1)
  
file_path: Quantus/tests/metrics/test_relative_stability.py
imports: from __future__ import annotations
  import pytest
  import torch
  from pytest_lazyfixture import lazy_fixture
  from typing import Dict, TYPE_CHECKING
  import functools
  from quantus.functions.perturb_func import *
  from quantus.functions.explanation_func import explain
  from quantus.metrics.robustness import RelativeInputStability, RelativeOutputStability, RelativeRepresentationStability
label_desc: Use image node_id 1 for calling a global function with example usage: predict(model, x_batch) and returns: model, model