diff --git a/src/pystencils_autodiff/backends/_torch_native.py b/src/pystencils_autodiff/backends/_torch_native.py index 48417af13d375c9c624d7b783bb207b437a07d5d..d43cb905e18b405a10b6048a9ee647e9f03fa597 100644 --- a/src/pystencils_autodiff/backends/_torch_native.py +++ b/src/pystencils_autodiff/backends/_torch_native.py @@ -62,7 +62,8 @@ def create_autograd_function(autodiff_obj, use_cuda): kwargs[field.name] = torch.zeros( field.shape, dtype=numpy_dtype_to_torch(field.dtype.numpy_dtype), - device=next(chain(args, kwargs.values())).device) + device='cuda' if use_cuda else 'cpu') # use device of tensor + output_tensors = OrderedDict({f.name: field_to_tensor_dict.get(f, kwargs[f.name]) for f in autodiff_obj.forward_output_fields}) diff --git a/src/pystencils_autodiff/backends/astnodes.py b/src/pystencils_autodiff/backends/astnodes.py index a040b060adbe957e0e092e2ebc0617a47b63a576..b8487f1395b8617eced9425a8d9302119187c179 100644 --- a/src/pystencils_autodiff/backends/astnodes.py +++ b/src/pystencils_autodiff/backends/astnodes.py @@ -55,7 +55,7 @@ class Header(JinjaCppFile): class TorchTensorDestructuring(DestructuringBindingsForFieldClass): CLASS_TO_MEMBER_DICT = { - FieldPointerSymbol: "data_ptr<{dtype}>()", + FieldPointerSymbol: "data<{dtype}>()", FieldShapeSymbol: "size({dim})", FieldStrideSymbol: "strides()[{dim}]" }