update includes

ef3a0025 · Natalia Gimelshein · 8124fba2 · ef3a0025 · ef3a0025
Commit ef3a0025 authored Oct 30, 2018 by Natalia Gimelshein
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

apex/optimizers/csrc/fused_adam_cuda.cpp apex/optimizers/csrc/fused_adam_cuda.cpp +1 -1

csrc/syncbn.cpp csrc/syncbn.cpp +1 -1

No files found.
--- a/apex/optimizers/csrc/fused_adam_cuda.cpp
+++ b/apex/optimizers/csrc/fused_adam_cuda.cpp
-#include <torch/torch.h>
+#include <torch/extension.h>

 // CUDA forward declaration
 void fused_adam_cuda(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode);

--- a/csrc/syncbn.cpp
+++ b/csrc/syncbn.cpp
-#include <torch/torch.h>
+#include <torch/extension.h>
 #include <ATen/ATen.h>

 #include <vector>