Disallow changing the device of a tensor via set_. (#18832)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18832
ghimport-source-id: fde4ad90541ba52dfa02bdd83466f17e6541e535

Stack from [ghstack](https://github.com/ezyang/ghstack):
* #18833 [STACK] Cache device on TensorImpl; clean up TensorImpl constructors.
* **#18832 [STACK] Disallow changing the device of a tensor via set_.**
* #18831 [STACK] Stop swapping in Storages of the wrong device for Tensors.

This is necessary to cache the device on a TensorImpl.

Differential Revision: D14766231

fbshipit-source-id: bba61634b2d6252ac0697b96033c9eea680956e8

1 files changed, 13 insertions, 2 deletions

diff --git a/aten/src/TH/THTensor.cpp b/aten/src/TH/THTensor.cpp
index bef67f9725..2be5ce7586 100644
--- a/aten/src/TH/THTensor.cpp
+++ b/aten/src/TH/THTensor.cpp
@@ -56,9 +56,10 @@ void THTensor_setStorageNd(THTensor *self, THStorage *storage, ptrdiff_t storage
   }
 
   /* storageOffset */
-  if(storageOffset < 0)
+  if(storageOffset < 0) {
     THError("Tensor: invalid storage offset");
-    self->set_storage_offset(storageOffset);
+  }
+  self->set_storage_offset(storageOffset);
 
   /* size and stride */
   THTensor_resizeNd(self, nDimension, size, stride);
@@ -160,5 +161,15 @@ void THTensor_stealAndSetStoragePtr(THTensor* tensor, THStorage* storage) {
   // Caffe2 might have tensors whose storages are null, but we
   // don't allow it in PyTorch.
   AT_ASSERT(storage);
+  // Caffe2 also has uninitialized dtype states, which we disallow here
+  AT_ASSERT(tensor->storage().dtype() == storage->dtype());
+
+  // We used to allow this, but this breaks device caching,
+  // see Note [We regret making Variable hold a Tensor]
+  // Let's put an actual error message for this one.
+  AT_CHECK(tensor->storage().device() == storage->device(),
+            "Attempted to set the storage of a tensor on device ", tensor->storage().device(),
+             " to a storage on different device ", storage->device(),
+            ".  This is no longer allowed; the devices must match.");
   tensor->set_storage(at::Storage(c10::intrusive_ptr<THStorage>::reclaim(storage)));
 }