Function getTensorMemoryByDevice

torch/utils/model_dump/code.js:536–549

(pickles)

}

function getTensorMemoryByDevice(pickles) {
  let all_tensors = [];
  for (const [name, pickle] of pickles) {
    const tensors = getTensorStorages(pickle);
    all_tensors.push(...tensors.values());
  }
  let result = {};
  for (const storage of all_tensors.values()) {
    const [dtype, key, device, numel] = storage;
    const size = computeTensorMemory(numel, dtype);
    result[device] = (result[device] || 0) + size;
  }
  return result;
}

// Make this a separate component so it is rendered lazily.
class OpenTensorMemorySection extends Component {