Hi,
I tried to use torch.utils.tensorboard (for replacing tensorboardX) and found that:
import os

import torch
import torch.nn.functional as F
from tqdm import tqdm
from torch.utils.data import DataLoader
os.environ['CUDA_VISIBLE_DEVICES'] = '4'
from torch.utils.tensorboard import SummaryWriter

(main code)

works normally, while
import os

import torch
import torch.nn.functional as F
from tqdm import tqdm
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
os.environ['CUDA_VISIBLE_DEVICES'] = '4'

(main code)

will still use gpu 0.
Am I using it wrong?
Thanks!