Update iris/src/world_model.py

iris/src/world_model.py

@@ -5,9 +5,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from models.kv_caching import KeysValues
-from models.slicer import Embedder, Head
-from models.transformer import Transformer
+from .models.kv_caching import KeysValues
+from .models.slicer import Embedder, Head
+from .models.transformer import Transformer
 
 class WorldModel(nn.Module):
     def __init__(self, obs_vocab_size: int, act_vocab_size: int, config: dict) -> None: