无法将gpu与google colab上的Pytorch一起使用

Question

我打开了一个谷歌协作笔记本，在上面运行一个python包，打算用GPU处理它

接下来link我选择了GPU选项（在Runtime选项中）并下载了所需的软件包，以便将GPU与Pytorch和Cuda一起使用。然而，出于某种原因，它表明存在一个CPU而不是GPU

安装软件包（需要使用conda）

!pip install -q condacolab

import condacolab
condacolab.install()
✨🍰✨ Everything looks OK!

!mamba install -c conda-forge scikit-learn 
!mamba install -q openmm
!mamba install -q   pytorch torchvision  cudatoolkit=11.1 -c pytorch 
!mamba install -c rdkit rdkit

!mamba install -q ipykernel 
!mamba install -q matplotlib

连接到google drive，其中包含脚本的.py文件为：

from google.colab import drive
drive.mount('/content/drive/')

Mounted at /content/drive/

但在运行此脚本时：

torch.cuda.is_available()
False

它不承认

运行下面的脚本时，它也无法识别任何gpu：

def try_gpu(i=0): 
    """Return gpu(i) if exists, otherwise return cpu()."""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus(): 
  """Return all available GPUs, or [cpu(),] if no GPU exists."""
  devices = [
      torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]
  return devices if devices else [torch.device('cpu')]

try_gpu(), try_gpu(10), try_all_gpus()

输出均为“cpu”：

(device(type='cpu'), device(type='cpu'), [device(type='cpu')])

但是，当检查是否有带有Tenseflow的GPU时，我发现有：

import tensorflow as tf
device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
    raise SystemError('GPU device not found')
print('Found GPU at: {}'.format(device_name))

Found GPU at: /device:GPU:0

我不确定是否相关，但也添加了以下内容：

!nvidia-smi

Mon Jun  7 15:02:40 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 465.27       Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   77C    P0    33W /  70W |    222MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

关于如何使Pytork在google colab笔记本上与GPU协同工作的任何想法/提示，我们将不胜感激