frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update Module.cpp for THC changes

Browse Source
This commit is contained in:
Sam Gross
2016-09-30 11:13:14 -07:00
parent 3e1c88e3e0
commit c20828478e
1 changed files with 3 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
torch/csrc/cuda/Module.cpp
View File

@ -7,9 +7,7 @@
#include "THCP.h"
THCState _state;
THCState *state = &_state;
THCDeviceAllocator deviceAllocator;
THCState *state;
////////////////////////////////////////////////////////////////////////////////
// Class pointer cache
@ -84,11 +82,6 @@ static bool THCPModule_assignStateless()
void THCPModule_setDevice(int device)
{
THCudaCheck(cudaSetDevice(device));
THCRandom_setGenerator(state, device);
// The stream is per device, so update the stream as well
THCState_setStream(state, device, THCState_getCurrentStreamIndex(state));
THCState_setBlasHandle(state, device, THCState_getCurrentBlasHandleIndex(state));
}
PyObject * THCPModule_setDevice_wrap(PyObject *self, PyObject *arg)
@ -213,7 +206,8 @@ PyObject * THCPModule_initialSeed(PyObject *_unused)
bool THCPModule_initCuda(PyObject *module_dict) {
#define ASSERT_TRUE(cond) if (!(cond)) { return false; }
THCCachingAllocator_init(&state->cudaDeviceAllocator);
state = THCState_alloc();
THCCachingAllocator_init(THCState_getDeviceAllocator(state));
THCudaInit(state);
#ifdef USE_MAGMA