Merge pull request #639 from AznamirWoW/caching_fix

Reverted the caching functionality to avoid running out of memory
IAHispano · Aug 27, 2024 · 9d39f54 · 9d39f54
2 parents 509ed68 + 1263d34
commit 9d39f54
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/rvc/train/train.py b/rvc/train/train.py
@@ -533,7 +533,7 @@ def train_and_evaluate(
     net_d.train()
 
     # Data caching
-    if True:
+    if cache_data_in_gpu:
         data_iterator = cache
         if cache == []:
             for batch_idx, info in enumerate(train_loader):
@@ -559,7 +559,7 @@ def train_and_evaluate(
                         wave_lengths,
                         sid,
                     ) = info
-                if cache_data_in_gpu == True and torch.cuda.is_available():
+                if torch.cuda.is_available():
                     phone = phone.cuda(rank, non_blocking=True)
                     phone_lengths = phone_lengths.cuda(rank, non_blocking=True)
                     if pitch_guidance == True:
@@ -604,6 +604,8 @@ def train_and_evaluate(
                     )
         else:
             shuffle(cache)
+    else:
+        data_iterator = enumerate(train_loader)
 
     epoch_recorder = EpochRecorder()
     with tqdm(total=len(train_loader), leave=False) as pbar: