Merge commit 'fd97d92479e32e550866adfd1f0465e4cfa5e581'

Author: soumith

torch/lib/ATen/TensorImpl.h

@@ -20,7 +20,7 @@ struct TensorImpl {
  virtual int64_t dim() = 0;
  virtual Scalar localScalar() = 0;
  virtual void assign_(Scalar s) = 0;
- virtual void * unsafeGetTH() = 0;
+ virtual void * unsafeGetTH(bool retain) = 0;
  void retain() {
  ++refcount;
  }

torch/lib/ATen/templates/Tensor.h

@@ -29,7 +29,7 @@ struct Tensor {
 
  Tensor()
  : pImpl(nullptr){}
- explicit Tensor(TensorImpl * self, bool retain = true)
+ explicit Tensor(TensorImpl * self, bool retain)
  : pImpl(self) {
  if(pImpl != nullptr && retain)
  pImpl->retain();
@@ -67,7 +67,7 @@ struct Tensor {
  Tensor().swap(*this);
  }
  void reset(TensorImpl * rhs) {
- Tensor(rhs).swap(*this);
+ Tensor(rhs,true).swap(*this);
  }
  void reset(TensorImpl * rhs, bool retain) {
  Tensor(rhs, retain).swap(*this );
@@ -126,8 +126,8 @@ struct Tensor {
  template<typename T>
  T * data() const;
 
- void * unsafeGetTH() {
- return pImpl->unsafeGetTH();
+ void * unsafeGetTH(bool retain) {
+ return pImpl->unsafeGetTH(retain);
  }
 
  //toLongData(), toFloatData() etc.

torch/lib/ATen/templates/TensorDerived.cpp

@@ -36,7 +36,9 @@ int64_t ${Tensor}::dim() {
 const char * ${Tensor}::typeString() {
  return "${Type}";
 }
-void * ${Tensor}::unsafeGetTH() {
+void * ${Tensor}::unsafeGetTH(bool retain) {
+ if (retain)
+ ${THTensor}_retain(${state,} tensor);
  return tensor;
 }
 

torch/lib/ATen/templates/TensorDerived.h

@@ -20,7 +20,7 @@ struct ${Tensor} : public TensorImpl {
  virtual int64_t dim() override;
  virtual Scalar localScalar() override;
  virtual void assign_(Scalar s) override;
- virtual void * unsafeGetTH() override;
+ virtual void * unsafeGetTH(bool retain) override;
  static const char * typeString();
 
 //TODO(zach): sort of friend permissions later so this

torch/lib/ATen/templates/Type.h

@@ -109,7 +109,7 @@ struct Type {
  virtual std::unique_ptr<Storage> storage(size_t size) = 0;
  virtual std::unique_ptr<Storage> storageFromBlob(void * data, int64_t size) = 0;
  virtual std::unique_ptr<Generator> generator() = 0;
- virtual Tensor unsafeTensorFromTH(void * th_pointer) = 0;
+ virtual Tensor unsafeTensorFromTH(void * th_pointer, bool retain) = 0;
  virtual const char * toString() const = 0;
  Type & toBackend(Backend b);
  Type & toScalarType(ScalarType s);

torch/lib/ATen/templates/TypeDerived.cpp

@@ -33,8 +33,10 @@ std::unique_ptr<Storage> ${Type}::storageFromBlob(void * data, int64_t size) {
  return std::unique_ptr<Storage>(
  new ${Storage}(context,data,size));
 }
-Tensor ${Type}::unsafeTensorFromTH(void * th_pointer) {
- return Tensor(new ${Tensor}(context,(${THTensor}*)(th_pointer)));
+Tensor ${Type}::unsafeTensorFromTH(void * th_pointer,bool retain) {
+ if (retain)
+ ${THTensor}_retain(${state,} (${THTensor}*) th_pointer);
+ return Tensor(new ${Tensor}(context,(${THTensor}*)(th_pointer)),false);
 }
 std::unique_ptr<Generator> ${Type}::generator() {
  return std::unique_ptr<Generator>(new ${Generator}(context));

torch/lib/ATen/templates/TypeDerived.h

@@ -20,7 +20,7 @@ struct ${Type} : public Type {
  virtual const char * toString() const override;
  virtual TypeID ID() const override;
  static const char * typeString();
- Tensor unsafeTensorFromTH(void * th_pointer) override;
+ Tensor unsafeTensorFromTH(void * th_pointer, bool retain) override;
 
  // example
  // virtual Tensor * add(Tensor & a, Tensor & b) override;

torch/lib/ATen/test/basic.cpp

@@ -200,7 +200,7 @@ static void test(Type & type) {
  int a = 4;
  THFloatTensor *t = THFloatTensor_newWithSize2d(a, a);
  THFloatTensor_fill(t, a);
- Tensor tt = CPU(kFloat).unsafeTensorFromTH(t);
+ Tensor tt = CPU(kFloat).unsafeTensorFromTH(t,false);
  std::cout << tt << std::endl;
  }
  {