Skip to content

Instantly share code, notes, and snippets.

@HeinrichAD

HeinrichAD/ProcessorUnit.py

Last active November 27, 2020 21:46
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save HeinrichAD/eea00d50c5702cbf63082983b6a2d763 to your computer and use it in GitHub Desktop.
Save HeinrichAD/eea00d50c5702cbf63082983b6a2d763 to your computer and use it in GitHub Desktop.
Download ZIP
Google Colab Processor Unit / Runtime Type Detector.
Display the source blob
Display the rendered blob
Raw
colab_pu_detector.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Raw
ProcessorUnit.py
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