Small fix in TensorflowConversion class (#121)
Signed-off-by: Alex Black <blacka101@gmail.com>master
parent
0a3e78b540
commit
b370544b8f
|
@ -296,7 +296,7 @@ public class TensorflowConversion {
|
||||||
* @return
|
* @return
|
||||||
*/
|
*/
|
||||||
public static String defaultDeviceForThread() {
|
public static String defaultDeviceForThread() {
|
||||||
Integer deviceForThread = Nd4j.getAffinityManager().getDeviceForThread(Thread.currentThread());
|
Integer deviceForThread = Nd4j.getAffinityManager().getDeviceForCurrentThread();
|
||||||
String deviceName = null;
|
String deviceName = null;
|
||||||
//gpu
|
//gpu
|
||||||
if(Nd4j.getBackend().getClass().getName().contains("JCublasBackend")) {
|
if(Nd4j.getBackend().getClass().getName().contains("JCublasBackend")) {
|
||||||
|
|
Loading…
Reference in New Issue