Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
IPython profile to automatically initialize PyCUDA on a specified device
#!/usr/bin/env python
"""
IPython 0.10 profile to automatically initialize PyCUDA on a specified
device.
"""
import sys
import atexit
import IPython.ipapi
ip = IPython.ipapi.get()
ip.ex("import numpy as np")
ip.ex("import atexit")
ip.ex("import pycuda.driver as drv")
ip.ex("drv.init()")
dev_num_str = sys.argv[-1] if sys.argv[-1].isdigit() else '0'
ip.ex("dev = drv.Device(%s)" % dev_num_str)
ip.ex("atexit.register(dev.make_context().pop)")
ip.ex("import pycuda.gpuarray as gpuarray")
ip.ex("print ")
ip.ex("print '*** pycuda initialized on ' + dev.name() + ' [%s] ***'" % dev_num_str)
ip.ex("print ")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.