Hmm maybe it would work if .cuda() just did nothing when the user had set the device to ‘cpu’. The following kind of works to give what I’m after.
from contextlib import contextmanager                                           
import torch                                                                          

torch.manually_set_device='none'                                                   
                                                                                     
@contextmanager                                                                    
def set_torch_device(new_device):                                                  
    prev_device = torch.manually_set_device                                        
    torch.manually_set_device = new_device                                         
    yield                                                                          
    torch.manually_set_device = prev_device                                        
                                                                                     
def maybe_cuda(t):                                                                 
    if torch.manually_set_device == 'none':                                        
        return t.cuda()                                                            
    else:                                                                          
        return t.to(torch.manually_set_device)                                     
                                                                                     
def make_tensor():                                                                 
    return maybe_cuda(torch.Tensor([1.,2.]))                                       
  
device = torch.device('cpu')                                                    
with set_torch_device(device):                                                                                                                                                                    
    t = make_tensor()                                                              
    print(t.device)                                                                
                                                                                     
t = make_tensor()                                                                  
print(t.device)