Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Google Colab Processor Unit / Runtime Type Detector.
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
from enum import Enum
from os import environ
class PU(Enum):
"""
Google Colab Processor Unit / Runtime Type Detector.
"""
CPU = 1
GPU = 2
TPU = 3
def isCPU():
return not (PU.isGPU() or PU.isTPU())
def isGPU():
return environ["COLAB_GPU"] == '1'
def isTPU():
return "TPU_NAME" in environ
def getPU():
return PU.GPU if PU.isGPU() else PU.TPU if PU.isTPU() else PU.CPU
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment