TestMemoryReportingInfo Class — pytorch Architecture

Source Code

aten/src/ATen/test/reportMemoryUsage.h lines 8–40

class TestMemoryReportingInfo : public c10::MemoryReportingInfoBase {
 public:
  struct Record {
    void* ptr;
    int64_t alloc_size;
    size_t total_allocated;
    size_t total_reserved;
    c10::Device device;
  };

  std::vector<Record> records;

  TestMemoryReportingInfo() = default;
  ~TestMemoryReportingInfo() override = default;

  void reportMemoryUsage(
      void* ptr,
      int64_t alloc_size,
      size_t total_allocated,
      size_t total_reserved,
      c10::Device device) override {
    records.emplace_back(
        Record{ptr, alloc_size, total_allocated, total_reserved, device});
  }

  bool memoryProfilingEnabled() const override {
    return true;
  }

  Record getLatestRecord() {
    return records.back();
  }
};