1 #pragma once 2 // TODO: Remove once torchvision has been updated to use the ATen header 3 #include <ATen/cuda/Atomic.cuh> 4