Add use_cuda arg to map_data
@@ -247,7 +247,7 @@ def irange(name, size, subsample_size=0, subsample=None, use_cuda=False):
yield i
-def map_data(name, data, fn, batch_size=0, batch_dim=0):
+def map_data(name, data, fn, batch_size=0, batch_dim=0, use_cuda=False):
"""
Data subsampling with the important property that all the data are conditionally independent.
@@ -259,15 +259,17 @@ def map_data(name, data, fn, batch_size=0, batch_dim=0):
:param callable fn: a function taking `(index, datum)` pairs, where `dataum = data[index]`
:param int batch_size: number of samples per batch, or zero for the entire dataset
:param int batch_dim: dimension to subsample for tensor inputs
+ :param bool use_cuda: Whether to use cuda tensors for `subsample` and `log_pdf`.
:return: a list of values returned by `fn`
"""
+ use_cuda = use_cuda or getattr(data, 'is_cuda', False)
if isinstance(data, (torch.Tensor, Variable)):
size = data.size(batch_dim)
- with iarange(name, size, batch_size, use_cuda=data.is_cuda) as batch:
+ with iarange(name, size, batch_size, use_cuda=use_cuda) as batch:
return fn(batch, data.index_select(batch_dim, batch))
else:
size = len(data)
- return [fn(i, data[i]) for i in irange(name, size, batch_size, use_cuda=data.is_cuda)]
+ return [fn(i, data[i]) for i in irange(name, size, batch_size, use_cuda=use_cuda)]
# XXX this should have the same call signature as torch.Tensor constructors
