text: import torch
  import numpy as np
  
  def compute_num_params(model, text=False):
      tot = int(sum([np.prod(p.shape) for p in model.parameters()]))
      if text:
          if tot >= 1e6:
              return '{:.1f}M'.format(tot / 1e6)
          else:
              return '{:.1f}K'.format(tot / 1e3)
      else:
          return tot
  
  def batched_predict(model, inp, coord, cell, bsize):
      with torch.no_grad():
          model.gen_feat(inp)
          n = coord.shape[1]
          ql = 0
          preds = []
          while ql < n:
              qr = min(ql + bsize, n)
              pred = model.query_rgb(coord[:, ql: qr, :], cell[:, ql: qr, :])
              preds.append(pred)
              ql = qr
          pred = torch.cat(preds, dim=1)
      return pred