Skip to content

Instantly share code, notes, and snippets.

@SomeoneSerge

SomeoneSerge/check-python.py

Created April 22, 2022 02:26
Show Gist options
  • Save SomeoneSerge/2b18573dba91eaa9bb4e9e70b914d655 to your computer and use it in GitHub Desktop.
Save SomeoneSerge/2b18573dba91eaa9bb4e9e70b914d655 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
check-python.py
#! /usr/bin/env nix-shell
#! nix-shell -i python -p "(import ./. { config = { allowUnfree = true; cudaSupport = true; }; }).python3.withPackages (ps: with ps; [ pytorch ])"
import torch
a = torch.randn(1, 3, 128, 128, device="cuda")
torch.nn.functional.conv2d(a, torch.randn(1, 3, 2, 2, device="cuda"))
print("ok")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment