Commit 1581d5a1 authored by wywywywy's avatar wywywywy
Browse files

Made device agnostic

parent 6df316c8

extensions-builtin/LDSR/ldsr_model_arch.py

+6 −3
Original line number Diff line number Diff line
@@ -110,6 +110,7 @@ class LDSR: 
        down_sample_method = 'Lanczos'



        gc.collect()

        if torch.cuda.is_available:

            torch.cuda.empty_cache()



        im_og = image
@@ -147,7 +148,9 @@ class LDSR: 


        del model

        gc.collect()

        if torch.cuda.is_available:

            torch.cuda.empty_cache()



        return a
@@ -162,7 +165,7 @@ def get_cond(selected_path): 
    c = rearrange(c, '1 c h w -> 1 h w c')

    c = 2. * c - 1.



    c = c.to(torch.device("cuda"))

    c = c.to(shared.device)

    example["LR_image"] = c

    example["image"] = c_up

CRA Git | Maintained and supported by SUSTech CRA and CCSE