text: import os
  import torch
  import random
  import numpy as np
  import math
  import torch.distributions as pyd
  import psutil
  import datetime
  
  
  class TanhTransform(pyd.transforms.Transform):
      domain = pyd.constraints.real
      codomain = pyd.constraints.interval(-1.0, 1.0)
      bijective = True
      sign = +1
  
      def __init__(self, cache_size=1):
          super().__init__(cache_size=cache_size)
  
      @staticmethod
      def atanh(x):
          return 0.5 * (x.log1p() - (-x).log1p())
  
      def __eq__(self, other):
          return isinstance(other, TanhTransform)
  
      def _call(self, x):
          return x.tanh()
  
      def _inverse(self, y):
          # We do not clamp to the boundary here as it may degrade the performance of certain algorithms.
          # one should use `cache_size=1` instead
          return self.atanh(y)
  
      def log_abs_det_jacobian(self, x, y):
          # We use a formula that is more numerically stable, see details in the following link
          # https://github.com/tensorflow/probability/commit/ef6bb176e0ebd1cf6e25c6b5cecdd2428c22963f#diff-e120f70e92e6741bca649f04fcd907b7
          return 2. * (math.log(2.) - x - F.softplus(-2. * x))
  
  
  class SquashedNormal(pyd.transformed_distribution.TransformedDistribution):
      def __init__(self, loc, scale):
          self.loc = loc
          self.scale = scale
  
          self.base_dist = pyd.Normal(loc, scale)
          transforms = [TanhTransform()]
          super().__init__(self.base_dist, transforms)
  
      @property
      def mean(self):
          mu = self.loc
          for tr in self.transforms:
              mu = tr(mu)
          return mu
  
  
  def soft_update_params(net, target_net, tau):
      for param, target_param in zip(net.parameters(), target_net.parameters()):
          target_param.data.copy_(tau * param.data +
                                  (1 - tau) * target_param.data)
  
  
  class eval_mode(object):
      def __init__(self, *models):
          self.models = models
  
      def __enter__(self):
          self.prev_states = []
          for model in self.models:
              self.prev_states.append(model.training)
              model.train(False)
  
      def __exit__(self, *args):
          for model, state in zip(self.models, self.prev_states):
              model.train(state)
          return False
  
  
  class train_mode(object):
      def __init__(self, *models):
          self.models = models
  
      def __enter__(self):
          self.prev_states = []
          for model in self.models:
              self.prev_states.append(model.training)
              model.train(True)
  
      def __exit__(self, *args):
          for model, state in zip(self.models, self.prev_states):
              model.train(state)
          return False
  
  
  def set_one_thread():
      '''
      Set number of threads for pytorch to 1
      '''
      os.environ['OMP_NUM_THREADS'] = '1'
      os.environ['MKL_NUM_THREADS'] = '1'
      torch.set_num_threads(1)
  
  
  def to_tensor(x, device):
      '''
      Convert an array to tensor
      '''
      x = torch.as_tensor(x, device=device, dtype=torch.float32)
      return x
  
  
  def to_numpy(t):
      '''
      Convert a tensor to numpy
      '''
      return t.cpu().detach().numpy()
  
  
  def set_random_seed(seed):
      '''
      Set all random seeds
      '''
      random.seed(seed)
      np.random.seed(seed)
      torch.manual_seed(seed)
      if torch.cuda.is_available():
          torch.cuda.manual_seed_all(seed)
          torch.backends.cudnn.deterministic = True
          torch.backends.cudnn.benchmark = False
  
  
  def make_dir(dir):
      if not os.path.exists(dir):
          os.makedirs(dir, exist_ok=True)
  
  
  def get_time_str():
      return datetime.datetime.now().strftime("%y.%m.%d-%H:%M:%S")
  
  
  def rss_memory_usage():
      '''
      Return the resident memory usage in MB
      '''
      process = psutil.Process(os.getpid())
      mem = process.memory_info().rss / float(2 ** 20)
      return mem
  