fix(//cpp/api): Better inital condition for the dataloader iterator to

address datarace issue Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
pytorch · Apr 24, 2020 · 8d22bdd · 8d22bdd
1 parent 5c0d737
commit 8d22bdd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/cpp/api/include/trtorch/ptq.h b/cpp/api/include/trtorch/ptq.h
@@ -27,7 +27,7 @@ class Int8Calibrator : Algorithm {
     using Batch = typename DataLoader::super::BatchType;
 public:
     Int8Calibrator(DataLoaderUniquePtr dataloader, const std::string& cache_file_path, bool use_cache)
-      : dataloader_(dataloader.get()), it_(dataloader_->begin()), cache_file_path_(cache_file_path), use_cache_(use_cache) {}
+      : dataloader_(dataloader.get()), it_(dataloader_->end()), cache_file_path_(cache_file_path), use_cache_(use_cache) {}
 
     int getBatchSize() const override {
         // HACK: TRTorch only uses explict batch sizing, INT8 Calibrator does not

diff --git a/cpp/ptq/main.cpp b/cpp/ptq/main.cpp
@@ -121,6 +121,13 @@ int main(int argc, const char* argv[]) {
     auto execution_timer = timers::PreciseCPUTimer();
     auto images = (*(*eval_dataloader).begin()).data.to(torch::kCUDA);
 
+    execution_timer.start();
+    mod.forward({images});
+    execution_timer.stop();
+    std::cout << "Latency of JIT model FP32 (Batch Size 32): " << execution_timer.milliseconds() << "ms" << std::endl;
+
+    execution_timer.reset();
+
     execution_timer.start();
     trt_mod.forward({images});
     execution_timer.stop();