Tracked new device argument in is_pinned/pin_memory for nested tensor

Reviewed By: cpuhrsch

Differential Revision: D29467785

fbshipit-source-id: 966c05ef16b7955afee982c168002c2ab7ed644a

Author: ezyang

nestedtensor/csrc/functions.cpp

@@ -159,9 +159,9 @@ Tensor NestedTensor__log_softmax(
       [&](Tensor a) { return at::_log_softmax(a, dim_, half_to_float); }, self);
 }
 
-Tensor NestedTensor_pin_memory(const Tensor& self) {
+Tensor NestedTensor_pin_memory(const Tensor& self, c10::optional<Device> device) {
   return map_nested_tensor(
-      [](Tensor tensor) { return at::native::pin_memory(tensor); }, self);
+      [&](Tensor tensor) { return at::native::pin_memory(tensor, device); }, self);
 }
 
 Tensor NestedTensor_flatten(

nestedtensor/csrc/nested_tensor_impl.cpp

@@ -135,7 +135,10 @@ Tensor NestedTensor_contiguous(const Tensor& self, MemoryFormat memory_format) {
       std::shared_ptr<NestedTensorStorage>(ps_base));
 }
 
-bool NestedTensor_is_pinned(const Tensor& self) {
+bool NestedTensor_is_pinned(const Tensor& self, c10::optional<Device> device) {
+  TORCH_CHECK(
+      !device.has_value() || device->is_cuda(),
+      "nested tensor doesn't support non-CUDA pinned memory");
   return get_nested_tensor_impl(self)->is_pinned();
 }