Function computeTensorMemory

torch/utils/model_dump/code.js:455–479

(numel, dtype)

}

function computeTensorMemory(numel, dtype) {
  const sizes = {
    "Byte": 1,
    "Char": 1,
    "Short": 2,
    "Int": 4,
    "Long": 8,
    "Half": 2,
    "Float": 4,
    "Double": 8,
    "ComplexHalf": 4,
    "ComplexFloat": 8,
    "ComplexDouble": 16,
    "Bool": 1,
    "QInt8": 1,
    "QUInt8": 1,
    "QInt32": 4,
    "BFloat16": 2,
  };
  let dtsize = sizes[dtype];
  if (!dtsize) {
    throw new Error("Unrecognized dtype: " + dtype);
  }
  return numel * dtsize;
}

// TODO: Maybe track by dtype as well.
// TODO: Maybe distinguish between visible size and storage size.