From 06abc8cb409f45b956ea60ea4a40f82303eeaaf0 Mon Sep 17 00:00:00 2001
From: Philippe Tillet <ptillet@g.harvard.edu>
Date: Tue, 3 Nov 2020 16:02:02 -0500
Subject: [PATCH] [GENERAL] Fix compatibility issue with older Torch versions

---
 python/src/launch.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/src/launch.cc b/python/src/launch.cc
index 5e698b096..6bf8edd5b 100644
--- a/python/src/launch.cc
+++ b/python/src/launch.cc
@@ -18,8 +18,8 @@ std::shared_ptr<drv::device> host_device;
 std::shared_ptr<drv::context> host_context;
 std::shared_ptr<drv::stream> host_stream;
 
-int64_t cdiv_sum(torch::Tensor& x, int64_t div){
-  TORCH_CHECK(x.unsafeGetTensorImpl()->key_set().has(c10::DispatchKey::CPU), "Argument of cdiv_sum must be a CPU tensor")
+int64_t cdiv_sum(torch::Tensor x, int64_t div){
+  TORCH_CHECK(!x.is_cuda(), "Argument of cdiv_sum must be a CPU tensor")
   auto _x  = x.accessor<int, 1>();
   int64_t ret = 0;
   for(size_t i = 0; i < x.size(0); i++)