Commit eb413124 authored by rusty1s's avatar rusty1s
Browse files

fixed gpu tests

parent c7f73ca2
......@@ -9,7 +9,7 @@ __global__ void gridKernel(int64_t *self, TensorInfo<T> posInfo, T *size,
KERNEL_LOOP(i, nNodes) {
T *pos = posInfo.data + i * posInfo.stride[0];
int64_t coef = 1, value = 0;
for (ptrdiff_t d = 0; d < posInfo.dims * posInfo.stride[1]; d += posInfo.stride[1]) {
for (ptrdiff_t d = 0; d < posInfo.size[1]; d += posInfo.stride[1]) {
value += coef * THCNumerics<T>::floor(THCNumerics<T>::div(pos[d], size[d]));
coef *= count[d];
}
......
......@@ -6,12 +6,13 @@ void THCTensor_(grid)(THCState *state, THCudaLongTensor *self, THCTensor *pos, T
THCudaLongTensor *count) {
THCAssertSameGPU(THCTensor_(checkGPU)(state, 4, self, pos, size, count));
ptrdiff_t nNodes = THCudaLongTensor_nElement(state, self);
int64_t *selfData = THCudaLongTensor_data(state, self);
TensorInfo<real> posInfo = THCTensor_(getTensorInfo)(state, pos);
real *sizeData = THCTensor_(data)(state, size);
int64_t *countData = THCudaLongTensor_data(state, count);
ptrdiff_t nNodes = THCudaLongTensor_nElement(state, self);
KERNEL_REAL_RUN(gridKernel, nNodes, selfData, posInfo, sizeData, countData);
}
......
......@@ -46,7 +46,7 @@ def test_grid_cluster_cpu(tensor, i):
@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
@pytest.mark.parametrize('tensor,i', product(tensors, range(len(tests))))
def test_grid_cluster_gpu(tensor, i):
def test_grid_cluster_gpu(tensor, i): # pragma: no cover
data = tests[i]
pos = getattr(torch.cuda, tensor)(data['pos'])
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment