code: def test_set_on_device(self, ae, train_dataset, training_config):
      trainer = BaseTrainer(
          model=ae,
          train_dataset=train_dataset,
          training_config=training_config,
      )
  
      device = "cuda" if torch.cuda.is_available() else "cpu"
  
      if training_config.no_cuda:
          assert next(trainer.model.parameters()).device == "cpu"
  
      else:
          next(trainer.model.parameters()).device == device
  
imports: import os
  from copy import deepcopy
  import pytest
  import torch
  from pythae.customexception import ModelError
  from pythae.models import AE, RHVAE, VAE, AEConfig, RHVAEConfig, VAEConfig
  from pythae.trainers import BaseTrainer, BaseTrainerConfig
  from tests.data.custom_architectures import *
label_desc: Use image node_id 7 for calling the Test_Device_Checks obj's underlying member method code with example usage: obj.test_set_on_device(ae, train_dataset, training_config) without return types