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