Skip to content

Instantly share code, notes, and snippets.

@ij96

ij96/torch_header.py

Created March 1, 2021 11:29
Show Gist options
  • Save ij96/f66d15e6d9b8ceaebe8eb2312e02a252 to your computer and use it in GitHub Desktop.
Save ij96/f66d15e6d9b8ceaebe8eb2312e02a252 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
torch_header.py
import torch
import torchvision
# for defining networks etc.
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms.functional as F_v
# specify torch device
TORCH_DEVICE = 'cuda:0'
# fallback to CPU if CUDA not available
TORCH_DEVICE = torch.device(TORCH_DEVICE) if torch.cuda.is_available() else torch.device('cpu')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment